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  }