feat: login 2fa

internal/user/otp.go

@@ -0,0 +1,39 @@
+package user
+
+import (
+	"bytes"
+	"crypto/sha1"
+	"encoding/hex"
+	"github.com/0xJacky/Nginx-UI/internal/crypto"
+	"github.com/0xJacky/Nginx-UI/model"
+	"github.com/pkg/errors"
+	"github.com/pquerna/otp/totp"
+)
+
+var (
+	ErrOTPCode      = errors.New("invalid otp code")
+	ErrRecoveryCode = errors.New("invalid recovery code")
+)
+
+func VerifyOTP(user *model.Auth, otp, recoveryCode string) (err error) {
+	if otp != "" {
+		decrypted, err := crypto.AesDecrypt(user.OTPSecret)
+		if err != nil {
+			return err
+		}
+
+		if ok := totp.Validate(otp, string(decrypted)); !ok {
+			return ErrOTPCode
+		}
+	} else {
+		recoverCode, err := hex.DecodeString(recoveryCode)
+		if err != nil {
+			return err
+		}
+		k := sha1.Sum(user.OTPSecret)
+		if !bytes.Equal(k[:], recoverCode) {
+			return ErrRecoveryCode
+		}
+	}
+	return
+}

internal/user/user.go

@@ -1,53 +1,84 @@
 package user
 
 import (
+	"github.com/0xJacky/Nginx-UI/internal/logger"
 	"github.com/0xJacky/Nginx-UI/model"
+	"github.com/0xJacky/Nginx-UI/query"
 	"github.com/0xJacky/Nginx-UI/settings"
 	"github.com/dgrijalva/jwt-go"
+	"github.com/pkg/errors"
+	"strings"
 	"time"
 )
 
+const ExpiredTime = 24 * time.Hour
+
 type JWTClaims struct {
-	Name string `json:"name"`
+	Name   string `json:"name"`
+	UserID int    `json:"user_id"`
 	jwt.StandardClaims
 }
 
-func GetUser(name string) (user model.Auth, err error) {
+func BuildCacheTokenKey(token string) string {
+	var sb strings.Builder
+	sb.WriteString("token:")
+	sb.WriteString(token)
+	return sb.String()
+}
+
+func GetUser(name string) (user *model.Auth, err error) {
 	db := model.UseDB()
-	err = db.Where("name", name).First(&user).Error
+	user = &model.Auth{}
+	err = db.Where("name", name).First(user).Error
 	if err != nil {
 		return
 	}
 	return
 }
 
-func DeleteToken(token string) error {
-	db := model.UseDB()
-	return db.Where("token", token).Delete(&model.AuthToken{}).Error
+func DeleteToken(token string) {
+	q := query.AuthToken
+	_, _ = q.Where(q.Token.Eq(token)).Delete()
 }
 
-func CheckToken(token string) int64 {
-	db := model.UseDB()
-	return db.Where("token", token).Find(&model.AuthToken{}).RowsAffected
+func GetTokenUser(token string) (*model.Auth, bool) {
+	q := query.AuthToken
+	authToken, err := q.Where(q.Token.Eq(token)).First()
+	if err != nil {
+		return nil, false
+	}
+
+	if authToken.ExpiredAt < time.Now().Unix() {
+		DeleteToken(token)
+		return nil, false
+	}
+
+	u := query.Auth
+	user, err := u.FirstByID(authToken.UserID)
+	return user, err == nil
 }
 
-func GenerateJWT(name string) (string, error) {
+func GenerateJWT(user *model.Auth) (string, error) {
 	claims := JWTClaims{
-		Name: name,
+		Name:   user.Name,
+		UserID: user.ID,
 		StandardClaims: jwt.StandardClaims{
-			ExpiresAt: time.Now().Add(24 * time.Hour).Unix(),
+			ExpiresAt: time.Now().Add(ExpiredTime).Unix(),
 		},
 	}
+
 	unsignedToken := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
 	signedToken, err := unsignedToken.SignedString([]byte(settings.ServerSettings.JwtSecret))
 	if err != nil {
 		return "", err
 	}
 
-	db := model.UseDB()
-	err = db.Create(&model.AuthToken{
-		Token: signedToken,
-	}).Error
+	q := query.AuthToken
+	err = q.Create(&model.AuthToken{
+		UserID:    user.ID,
+		Token:     signedToken,
+		ExpiredAt: time.Now().Add(ExpiredTime).Unix(),
+	})
 
 	if err != nil {
 		return "", err
@@ -55,3 +86,50 @@ func GenerateJWT(name string) (string, error) {
 
 	return signedToken, err
 }
+
+func ValidateJWT(token string) (claims *JWTClaims, err error) {
+	if token == "" {
+		err = errors.New("token is empty")
+		return
+	}
+	unsignedToken, err := jwt.ParseWithClaims(
+		token,
+		&JWTClaims{},
+		func(token *jwt.Token) (interface{}, error) {
+			return []byte(settings.ServerSettings.JwtSecret), nil
+		},
+	)
+	if err != nil {
+		err = errors.New("parse with claims error")
+		return
+	}
+	claims, ok := unsignedToken.Claims.(*JWTClaims)
+	if !ok {
+		err = errors.New("convert to jwt claims error")
+		return
+	}
+	if claims.ExpiresAt < time.Now().UTC().Unix() {
+		err = errors.New("jwt is expired")
+	}
+	return
+}
+
+func CurrentUser(token string) (u *model.Auth, err error) {
+	// validate token
+	var claims *JWTClaims
+	claims, err = ValidateJWT(token)
+	if err != nil {
+		return
+	}
+
+	// get user by id
+	user := query.Auth
+	u, err = user.FirstByID(claims.UserID)
+	if err != nil {
+		return
+	}
+
+	logger.Info("[Current User]", u.Name)
+
+	return
+}