🎣 Open-Source Phishing Toolkit
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

102 lines
2.8 KiB

  1. package models
  2. import (
  3. "github.com/jinzhu/gorm"
  4. "gopkg.in/check.v1"
  5. )
  6. func (s *ModelsSuite) TestGetUserExists(c *check.C) {
  7. u, err := GetUser(1)
  8. c.Assert(err, check.Equals, nil)
  9. c.Assert(u.Username, check.Equals, "admin")
  10. }
  11. func (s *ModelsSuite) TestGetUserByUsernameWithExistingUser(c *check.C) {
  12. u, err := GetUserByUsername("admin")
  13. c.Assert(err, check.Equals, nil)
  14. c.Assert(u.Username, check.Equals, "admin")
  15. }
  16. func (s *ModelsSuite) TestGetUserDoesNotExist(c *check.C) {
  17. u, err := GetUser(100)
  18. c.Assert(err, check.Equals, gorm.ErrRecordNotFound)
  19. c.Assert(u.Username, check.Equals, "")
  20. }
  21. func (s *ModelsSuite) TestGetUserByAPIKeyWithExistingAPIKey(c *check.C) {
  22. u, err := GetUser(1)
  23. c.Assert(err, check.Equals, nil)
  24. u, err = GetUserByAPIKey(u.ApiKey)
  25. }
  26. func (s *ModelsSuite) TestGetUserByAPIKeyWithNotExistingAPIKey(c *check.C) {
  27. u, err := GetUser(1)
  28. c.Assert(err, check.Equals, nil)
  29. u, err = GetUserByAPIKey(u.ApiKey + "test")
  30. c.Assert(err, check.Equals, gorm.ErrRecordNotFound)
  31. c.Assert(u.Username, check.Equals, "")
  32. }
  33. func (s *ModelsSuite) TestGetUserByUsernameWithNotExistingUser(c *check.C) {
  34. u, err := GetUserByUsername("test user does not exist")
  35. c.Assert(err, check.Equals, gorm.ErrRecordNotFound)
  36. c.Assert(u.Username, check.Equals, "")
  37. }
  38. func (s *ModelsSuite) TestPutUser(c *check.C) {
  39. u, err := GetUser(1)
  40. u.Username = "admin_changed"
  41. err = PutUser(&u)
  42. c.Assert(err, check.Equals, nil)
  43. u, err = GetUser(1)
  44. c.Assert(u.Username, check.Equals, "admin_changed")
  45. }
  46. func (s *ModelsSuite) TestGeneratedAPIKey(c *check.C) {
  47. u, err := GetUser(1)
  48. c.Assert(err, check.Equals, nil)
  49. c.Assert(u.ApiKey, check.Not(check.Equals), "12345678901234567890123456789012")
  50. }
  51. func (s *ModelsSuite) verifyRoleCount(c *check.C, roleID, expected int64) {
  52. var adminCount int64
  53. err := db.Model(&User{}).Where("role_id=?", roleID).Count(&adminCount).Error
  54. c.Assert(err, check.Equals, nil)
  55. c.Assert(adminCount, check.Equals, expected)
  56. }
  57. func (s *ModelsSuite) TestDeleteLastAdmin(c *check.C) {
  58. // Create a new admin user
  59. role, err := GetRoleBySlug(RoleAdmin)
  60. c.Assert(err, check.Equals, nil)
  61. newAdmin := User{
  62. Username: "new-admin",
  63. Hash: "123456",
  64. ApiKey: "123456",
  65. Role: role,
  66. RoleID: role.ID,
  67. }
  68. err = PutUser(&newAdmin)
  69. c.Assert(err, check.Equals, nil)
  70. // Ensure that there are two admins
  71. s.verifyRoleCount(c, role.ID, 2)
  72. // Delete the newly created admin - this should work since we have more
  73. // than one current admin.
  74. err = DeleteUser(newAdmin.Id)
  75. c.Assert(err, check.Equals, nil)
  76. // Verify that we now have one admin
  77. s.verifyRoleCount(c, role.ID, 1)
  78. // Try to delete the last admin - this should fail since we always want at
  79. // least one admin active in Gophish.
  80. err = DeleteUser(1)
  81. c.Assert(err, check.Equals, ErrModifyingOnlyAdmin)
  82. // Verify that the admin wasn't deleted
  83. s.verifyRoleCount(c, role.ID, 1)
  84. }