crawlab/core/user/service_v2.go

package user

import (
	mongo2 "github.com/crawlab-team/crawlab-db/mongo"
	"github.com/crawlab-team/crawlab/core/constants"
	"github.com/crawlab-team/crawlab/core/errors"
	"github.com/crawlab-team/crawlab/core/interfaces"
	"github.com/crawlab-team/crawlab/core/models/models"
	"github.com/crawlab-team/crawlab/core/models/service"
	"github.com/crawlab-team/crawlab/core/utils"
	"github.com/crawlab-team/go-trace"
	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"time"
)

type ServiceV2 struct {
	jwtSecret        string
	jwtSigningMethod jwt.SigningMethod
	modelSvc         *service.ModelServiceV2[models.UserV2]
}

func (svc *ServiceV2) Init() (err error) {
	_, err = svc.modelSvc.GetOne(bson.M{"username": constants.DefaultAdminUsername}, nil)
	if err == nil {
		return nil
	}
	if err.Error() != mongo.ErrNoDocuments.Error() {
		return err
	}
	return svc.Create(
		constants.DefaultAdminUsername,
		constants.DefaultAdminPassword,
		constants.RoleAdmin,
		"",
		primitive.NilObjectID,
	)
}

func (svc *ServiceV2) SetJwtSecret(secret string) {
	svc.jwtSecret = secret
}

func (svc *ServiceV2) SetJwtSigningMethod(method jwt.SigningMethod) {
	svc.jwtSigningMethod = method
}

func (svc *ServiceV2) Create(username, password, role, email string, by primitive.ObjectID) (err error) {
	// validate options
	if username == "" || password == "" {
		return trace.TraceError(errors.ErrorUserMissingRequiredFields)
	}
	if len(password) < 5 {
		return trace.TraceError(errors.ErrorUserInvalidPassword)
	}

	// normalize options
	if role == "" {
		role = constants.RoleNormal
	}

	// check if user exists
	if u, err := svc.modelSvc.GetOne(bson.M{"username": username}, nil); err == nil && u != nil && !u.Id.IsZero() {
		return trace.TraceError(errors.ErrorUserAlreadyExists)
	}

	// transaction
	return mongo2.RunTransaction(func(ctx mongo.SessionContext) error {
		// add user
		u := models.UserV2{
			Username: username,
			Role:     role,
			Password: utils.EncryptMd5(password),
			Email:    email,
		}
		u.SetCreated(by)
		u.SetUpdated(by)
		_, err = svc.modelSvc.InsertOne(u)

		return err
	})
}

func (svc *ServiceV2) Login(username, password string) (token string, u *models.UserV2, err error) {
	u, err = svc.modelSvc.GetOne(bson.M{"username": username}, nil)
	if err != nil {
		return "", nil, err
	}
	if u.Password != utils.EncryptMd5(password) {
		return "", nil, errors.ErrorUserMismatch
	}
	token, err = svc.makeToken(u)
	if err != nil {
		return "", nil, err
	}
	return token, u, nil
}

func (svc *ServiceV2) CheckToken(tokenStr string) (u *models.UserV2, err error) {
	return svc.checkToken(tokenStr)
}

func (svc *ServiceV2) ChangePassword(id primitive.ObjectID, password string, by primitive.ObjectID) (err error) {
	u, err := svc.modelSvc.GetById(id)
	if err != nil {
		return err
	}
	u.Password = utils.EncryptMd5(password)
	u.SetCreatedBy(by)
	return svc.modelSvc.ReplaceById(id, *u)
}

func (svc *ServiceV2) MakeToken(user *models.UserV2) (tokenStr string, err error) {
	return svc.makeToken(user)
}

func (svc *ServiceV2) GetCurrentUser(c *gin.Context) (user interfaces.User, err error) {
	// token string
	tokenStr := c.GetHeader("Authorization")

	// user
	u, err := userSvc.CheckToken(tokenStr)
	if err != nil {
		return nil, err
	}

	return u, nil
}

func (svc *ServiceV2) makeToken(user *models.UserV2) (tokenStr string, err error) {
	token := jwt.NewWithClaims(svc.jwtSigningMethod, jwt.MapClaims{
		"id":       user.Id,
		"username": user.Username,
		"nbf":      time.Now().Unix(),
	})
	return token.SignedString([]byte(svc.jwtSecret))
}

func (svc *ServiceV2) checkToken(tokenStr string) (user *models.UserV2, err error) {
	token, err := jwt.Parse(tokenStr, svc.getSecretFunc())
	if err != nil {
		return
	}

	claim, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		err = errors.ErrorUserInvalidType
		return
	}

	if !token.Valid {
		err = errors.ErrorUserInvalidToken
		return
	}

	id, err := primitive.ObjectIDFromHex(claim["id"].(string))
	if err != nil {
		return user, err
	}
	username := claim["username"].(string)
	user, err = svc.modelSvc.GetById(id)
	if err != nil {
		err = errors.ErrorUserNotExists
		return
	}

	if username != user.Username {
		err = errors.ErrorUserMismatch
		return
	}

	return
}

func (svc *ServiceV2) getSecretFunc() jwt.Keyfunc {
	return func(token *jwt.Token) (interface{}, error) {
		return []byte(svc.jwtSecret), nil
	}
}

func NewUserServiceV2() (svc *ServiceV2, err error) {
	// service
	svc = &ServiceV2{
		modelSvc:         service.NewModelServiceV2[models.UserV2](),
		jwtSecret:        "crawlab",
		jwtSigningMethod: jwt.SigningMethodHS256,
	}

	// initialize
	if err := svc.Init(); err != nil {
		return nil, trace.TraceError(err)
	}

	return svc, nil
}

var userSvcV2 *ServiceV2

func GetUserServiceV2() (svc *ServiceV2, err error) {
	if userSvcV2 != nil {
		return userSvcV2, nil
	}
	svc, err = NewUserServiceV2()
	if err != nil {
		return nil, err
	}
	userSvcV2 = svc
	return svc, nil
}