github.com/Ne0nd0g/gophish@v0.7.1-0.20190220040016-11493024a07d/middleware/middleware_test.go (about) 1 package middleware 2 3 import ( 4 "net/http" 5 "net/http/httptest" 6 "testing" 7 8 "github.com/gophish/gophish/config" 9 ctx "github.com/gophish/gophish/context" 10 "github.com/gophish/gophish/models" 11 "github.com/stretchr/testify/suite" 12 ) 13 14 var successHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 15 w.Write([]byte("success")) 16 }) 17 18 type MiddlewareSuite struct { 19 suite.Suite 20 } 21 22 func (s *MiddlewareSuite) SetupSuite() { 23 conf := &config.Config{ 24 DBName: "sqlite3", 25 DBPath: ":memory:", 26 MigrationsPath: "../db/db_sqlite3/migrations/", 27 } 28 err := models.Setup(conf) 29 if err != nil { 30 s.T().Fatalf("Failed creating database: %v", err) 31 } 32 } 33 34 // MiddlewarePermissionTest maps an expected HTTP Method to an expected HTTP 35 // status code 36 type MiddlewarePermissionTest map[string]int 37 38 // TestEnforceViewOnly ensures that only users with the ModifyObjects 39 // permission have the ability to send non-GET requests. 40 func (s *MiddlewareSuite) TestEnforceViewOnly() { 41 permissionTests := map[string]MiddlewarePermissionTest{ 42 models.RoleAdmin: MiddlewarePermissionTest{ 43 http.MethodGet: http.StatusOK, 44 http.MethodHead: http.StatusOK, 45 http.MethodOptions: http.StatusOK, 46 http.MethodPost: http.StatusOK, 47 http.MethodPut: http.StatusOK, 48 http.MethodDelete: http.StatusOK, 49 }, 50 models.RoleUser: MiddlewarePermissionTest{ 51 http.MethodGet: http.StatusOK, 52 http.MethodHead: http.StatusOK, 53 http.MethodOptions: http.StatusOK, 54 http.MethodPost: http.StatusOK, 55 http.MethodPut: http.StatusOK, 56 http.MethodDelete: http.StatusOK, 57 }, 58 } 59 for r, checks := range permissionTests { 60 role, err := models.GetRoleBySlug(r) 61 s.Nil(err) 62 63 for method, expected := range checks { 64 req := httptest.NewRequest(method, "/", nil) 65 response := httptest.NewRecorder() 66 67 req = ctx.Set(req, "user", models.User{ 68 Role: role, 69 RoleID: role.ID, 70 }) 71 72 EnforceViewOnly(successHandler).ServeHTTP(response, req) 73 s.Equal(response.Code, expected) 74 } 75 } 76 } 77 78 func (s *MiddlewareSuite) TestRequirePermission() { 79 middleware := RequirePermission(models.PermissionModifySystem) 80 handler := middleware(successHandler) 81 82 permissionTests := map[string]int{ 83 models.RoleUser: http.StatusForbidden, 84 models.RoleAdmin: http.StatusOK, 85 } 86 87 for role, expected := range permissionTests { 88 req := httptest.NewRequest(http.MethodGet, "/", nil) 89 response := httptest.NewRecorder() 90 // Test that with the requested permission, the request succeeds 91 role, err := models.GetRoleBySlug(role) 92 s.Nil(err) 93 req = ctx.Set(req, "user", models.User{ 94 Role: role, 95 RoleID: role.ID, 96 }) 97 handler.ServeHTTP(response, req) 98 s.Equal(response.Code, expected) 99 } 100 } 101 102 func TestMiddlewareSuite(t *testing.T) { 103 suite.Run(t, new(MiddlewareSuite)) 104 }