github.com/glennzw/gophish@v0.8.1-0.20190824020715-24fe998a3aa0/middleware/middleware_test.go (about)

     1  package middleware
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"testing"
     8  
     9  	"github.com/gophish/gophish/config"
    10  	ctx "github.com/gophish/gophish/context"
    11  	"github.com/gophish/gophish/models"
    12  	"github.com/stretchr/testify/suite"
    13  )
    14  
    15  var successHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    16  	w.Write([]byte("success"))
    17  })
    18  
    19  type MiddlewareSuite struct {
    20  	suite.Suite
    21  	apiKey string
    22  }
    23  
    24  func (s *MiddlewareSuite) SetupSuite() {
    25  	conf := &config.Config{
    26  		DBName:         "sqlite3",
    27  		DBPath:         ":memory:",
    28  		MigrationsPath: "../db/db_sqlite3/migrations/",
    29  	}
    30  	err := models.Setup(conf)
    31  	if err != nil {
    32  		s.T().Fatalf("Failed creating database: %v", err)
    33  	}
    34  	// Get the API key to use for these tests
    35  	u, err := models.GetUser(1)
    36  	s.Nil(err)
    37  	s.apiKey = u.ApiKey
    38  }
    39  
    40  // MiddlewarePermissionTest maps an expected HTTP Method to an expected HTTP
    41  // status code
    42  type MiddlewarePermissionTest map[string]int
    43  
    44  // TestEnforceViewOnly ensures that only users with the ModifyObjects
    45  // permission have the ability to send non-GET requests.
    46  func (s *MiddlewareSuite) TestEnforceViewOnly() {
    47  	permissionTests := map[string]MiddlewarePermissionTest{
    48  		models.RoleAdmin: MiddlewarePermissionTest{
    49  			http.MethodGet:     http.StatusOK,
    50  			http.MethodHead:    http.StatusOK,
    51  			http.MethodOptions: http.StatusOK,
    52  			http.MethodPost:    http.StatusOK,
    53  			http.MethodPut:     http.StatusOK,
    54  			http.MethodDelete:  http.StatusOK,
    55  		},
    56  		models.RoleUser: MiddlewarePermissionTest{
    57  			http.MethodGet:     http.StatusOK,
    58  			http.MethodHead:    http.StatusOK,
    59  			http.MethodOptions: http.StatusOK,
    60  			http.MethodPost:    http.StatusOK,
    61  			http.MethodPut:     http.StatusOK,
    62  			http.MethodDelete:  http.StatusOK,
    63  		},
    64  	}
    65  	for r, checks := range permissionTests {
    66  		role, err := models.GetRoleBySlug(r)
    67  		s.Nil(err)
    68  
    69  		for method, expected := range checks {
    70  			req := httptest.NewRequest(method, "/", nil)
    71  			response := httptest.NewRecorder()
    72  
    73  			req = ctx.Set(req, "user", models.User{
    74  				Role:   role,
    75  				RoleID: role.ID,
    76  			})
    77  
    78  			EnforceViewOnly(successHandler).ServeHTTP(response, req)
    79  			s.Equal(response.Code, expected)
    80  		}
    81  	}
    82  }
    83  
    84  func (s *MiddlewareSuite) TestRequirePermission() {
    85  	middleware := RequirePermission(models.PermissionModifySystem)
    86  	handler := middleware(successHandler)
    87  
    88  	permissionTests := map[string]int{
    89  		models.RoleUser:  http.StatusForbidden,
    90  		models.RoleAdmin: http.StatusOK,
    91  	}
    92  
    93  	for role, expected := range permissionTests {
    94  		req := httptest.NewRequest(http.MethodGet, "/", nil)
    95  		response := httptest.NewRecorder()
    96  		// Test that with the requested permission, the request succeeds
    97  		role, err := models.GetRoleBySlug(role)
    98  		s.Nil(err)
    99  		req = ctx.Set(req, "user", models.User{
   100  			Role:   role,
   101  			RoleID: role.ID,
   102  		})
   103  		handler.ServeHTTP(response, req)
   104  		s.Equal(response.Code, expected)
   105  	}
   106  }
   107  
   108  func (s *MiddlewareSuite) TestRequireAPIKey() {
   109  	req := httptest.NewRequest(http.MethodGet, "/", nil)
   110  	req.Header.Set("Content-Type", "application/json")
   111  	response := httptest.NewRecorder()
   112  	// Test that making a request without an API key is denied
   113  	RequireAPIKey(successHandler).ServeHTTP(response, req)
   114  	s.Equal(response.Code, http.StatusUnauthorized)
   115  }
   116  
   117  func (s *MiddlewareSuite) TestInvalidAPIKey() {
   118  	req := httptest.NewRequest(http.MethodGet, "/", nil)
   119  	query := req.URL.Query()
   120  	query.Set("api_key", "bogus-api-key")
   121  	req.URL.RawQuery = query.Encode()
   122  	req.Header.Set("Content-Type", "application/json")
   123  	response := httptest.NewRecorder()
   124  	RequireAPIKey(successHandler).ServeHTTP(response, req)
   125  	s.Equal(response.Code, http.StatusUnauthorized)
   126  }
   127  
   128  func (s *MiddlewareSuite) TestBearerToken() {
   129  	req := httptest.NewRequest(http.MethodGet, "/", nil)
   130  	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", s.apiKey))
   131  	req.Header.Set("Content-Type", "application/json")
   132  	response := httptest.NewRecorder()
   133  	RequireAPIKey(successHandler).ServeHTTP(response, req)
   134  	s.Equal(response.Code, http.StatusOK)
   135  }
   136  
   137  func TestMiddlewareSuite(t *testing.T) {
   138  	suite.Run(t, new(MiddlewareSuite))
   139  }