gophish/models/user_test.go

package models

import (
	"github.com/jinzhu/gorm"
	"gopkg.in/check.v1"
)

func (s *ModelsSuite) TestGetUserExists(c *check.C) {
	u, err := GetUser(1)
	c.Assert(err, check.Equals, nil)
	c.Assert(u.Username, check.Equals, "admin")
}

func (s *ModelsSuite) TestGetUserByUsernameWithExistingUser(c *check.C) {
	u, err := GetUserByUsername("admin")
	c.Assert(err, check.Equals, nil)
	c.Assert(u.Username, check.Equals, "admin")
}

func (s *ModelsSuite) TestGetUserDoesNotExist(c *check.C) {
	u, err := GetUser(100)
	c.Assert(err, check.Equals, gorm.ErrRecordNotFound)
	c.Assert(u.Username, check.Equals, "")
}

func (s *ModelsSuite) TestGetUserByAPIKeyWithExistingAPIKey(c *check.C) {
	u, err := GetUser(1)
	c.Assert(err, check.Equals, nil)

	got, err := GetUserByAPIKey(u.ApiKey)
	c.Assert(err, check.Equals, nil)
	c.Assert(got.Id, check.Equals, u.Id)
}

func (s *ModelsSuite) TestGetUserByAPIKeyWithNotExistingAPIKey(c *check.C) {
	u, err := GetUser(1)
	c.Assert(err, check.Equals, nil)

	u, err = GetUserByAPIKey(u.ApiKey + "test")
	c.Assert(err, check.Equals, gorm.ErrRecordNotFound)
	c.Assert(u.Username, check.Equals, "")
}

func (s *ModelsSuite) TestGetUserByUsernameWithNotExistingUser(c *check.C) {
	u, err := GetUserByUsername("test user does not exist")
	c.Assert(err, check.Equals, gorm.ErrRecordNotFound)
	c.Assert(u.Username, check.Equals, "")
}

func (s *ModelsSuite) TestPutUser(c *check.C) {
	u, _ := GetUser(1)
	u.Username = "admin_changed"
	err := PutUser(&u)
	c.Assert(err, check.Equals, nil)
	u, err = GetUser(1)
	c.Assert(err, check.Equals, nil)
	c.Assert(u.Username, check.Equals, "admin_changed")
}

func (s *ModelsSuite) TestGeneratedAPIKey(c *check.C) {
	u, err := GetUser(1)
	c.Assert(err, check.Equals, nil)
	c.Assert(u.ApiKey, check.Not(check.Equals), "12345678901234567890123456789012")
}

func (s *ModelsSuite) verifyRoleCount(c *check.C, roleID, expected int64) {
	var adminCount int64
	err := db.Model(&User{}).Where("role_id=?", roleID).Count(&adminCount).Error
	c.Assert(err, check.Equals, nil)
	c.Assert(adminCount, check.Equals, expected)
}

func (s *ModelsSuite) TestDeleteLastAdmin(c *check.C) {
	// Create a new admin user
	role, err := GetRoleBySlug(RoleAdmin)
	c.Assert(err, check.Equals, nil)
	newAdmin := User{
		Username: "new-admin",
		Hash:     "123456",
		ApiKey:   "123456",
		Role:     role,
		RoleID:   role.ID,
	}
	err = PutUser(&newAdmin)
	c.Assert(err, check.Equals, nil)

	// Ensure that there are two admins
	s.verifyRoleCount(c, role.ID, 2)

	// Delete the newly created admin - this should work since we have more
	// than one current admin.
	err = DeleteUser(newAdmin.Id)
	c.Assert(err, check.Equals, nil)

	// Verify that we now have one admin
	s.verifyRoleCount(c, role.ID, 1)

	// Try to delete the last admin - this should fail since we always want at
	// least one admin active in Gophish.
	err = DeleteUser(1)
	c.Assert(err, check.Equals, ErrModifyingOnlyAdmin)

	// Verify that the admin wasn't deleted
	s.verifyRoleCount(c, role.ID, 1)
}