2 files changed, 39 insertions, 61 deletions

diff --git a/web/service/setting.go b/web/service/setting.go
index 5443760c..62d66c11 100644
--- a/web/service/setting.go
+++ b/web/service/setting.go
@@ -48,7 +48,8 @@ var defaultValueMap = map[string]string{
 	"tgBotLoginNotify":            "true",
 	"tgCpu":                       "80",
 	"tgLang":                      "en-US",
-	"secretEnable":                "false",
+	"twoFactorEnable":             "false",
+	"twoFactorToken":              "",
 	"subEnable":                   "false",
 	"subTitle":                    "",
 	"subListen":                   "",
@@ -166,8 +167,7 @@ func (s *SettingService) ResetSettings() error {
 		return err
 	}
 	return db.Model(model.User{}).
-		Where("1 = 1").
-		Update("login_secret", "").Error
+		Where("1 = 1").Error
 }
 
 func (s *SettingService) getSetting(key string) (*model.Setting, error) {
@@ -318,6 +318,14 @@ func (s *SettingService) GetTgLang() (string, error) {
 	return s.getString("tgLang")
 }
 
+func (s *SettingService) GetTwoFactorEnable() (bool, error) {
+	return s.getBool("twoFactorEnable")
+}
+
+func (s *SettingService) GetTwoFactorToken() (string, error) {
+	return s.getString("twoFactorToken")
+}
+
 func (s *SettingService) GetPort() (int, error) {
 	return s.getInt("webPort")
 }
@@ -358,14 +366,6 @@ func (s *SettingService) GetRemarkModel() (string, error) {
 	return s.getString("remarkModel")
 }
 
-func (s *SettingService) GetSecretStatus() (bool, error) {
-	return s.getBool("secretEnable")
-}
-
-func (s *SettingService) SetSecretStatus(value bool) error {
-	return s.setBool("secretEnable", value)
-}
-
 func (s *SettingService) GetSecret() ([]byte, error) {
 	secret, err := s.getString("secret")
 	if secret == defaultValueMap["secret"] {
diff --git a/web/service/user.go b/web/service/user.go
index 72ae25a2..f0b04f52 100644
--- a/web/service/user.go
+++ b/web/service/user.go
@@ -8,10 +8,13 @@ import (
 	"x-ui/logger"
 	"x-ui/util/crypto"
 
+	"github.com/xlzd/gotp"
 	"gorm.io/gorm"
 )
 
-type UserService struct{}
+type UserService struct {
+	settingService SettingService
+}
 
 func (s *UserService) GetFirstUser() (*model.User, error) {
 	db := database.GetDB()
@@ -26,13 +29,13 @@ func (s *UserService) GetFirstUser() (*model.User, error) {
 	return user, nil
 }
 
-func (s *UserService) CheckUser(username string, password string, secret string) *model.User {
+func (s *UserService) CheckUser(username string, password string, twoFactorCode string) *model.User {
 	db := database.GetDB()
 
 	user := &model.User{}
 
 	err := db.Model(model.User{}).
-		Where("username = ? and login_secret = ?", username, secret).
+		Where("username = ?", username).
 		First(user).
 		Error
 	if err == gorm.ErrRecordNotFound {
@@ -42,69 +45,44 @@ func (s *UserService) CheckUser(username string, password string, secret string)
 		return nil
 	}
 
-	if crypto.CheckPasswordHash(user.Password, password) {
-		return user
+	if !crypto.CheckPasswordHash(user.Password, password) {
+		return nil
 	}
 
-	return nil
-}
-
-func (s *UserService) UpdateUser(id int, username string, password string) error {
-	db := database.GetDB()
-	hashedPassword, err := crypto.HashPasswordAsBcrypt(password)
-
+	twoFactorEnable, err := s.settingService.GetTwoFactorEnable()
 	if err != nil {
-		return err
+		logger.Warning("check two factor err:", err)
+		return nil
 	}
 
-	return db.Model(model.User{}).
-		Where("id = ?", id).
-		Updates(map[string]any{"username": username, "password": hashedPassword}).
-		Error
-}
+	if twoFactorEnable {
+		twoFactorToken, err := s.settingService.GetTwoFactorToken()
 
-func (s *UserService) UpdateUserSecret(id int, secret string) error {
-	db := database.GetDB()
-	return db.Model(model.User{}).
-		Where("id = ?", id).
-		Update("login_secret", secret).
-		Error
-}
+		if err != nil {
+			logger.Warning("check two factor token err:", err)
+			return nil
+		}
 
-func (s *UserService) RemoveUserSecret() error {
-	db := database.GetDB()
-	return db.Model(model.User{}).
-		Where("1 = 1").
-		Update("login_secret", "").
-		Error
-}
-
-func (s *UserService) GetUserSecret(id int) *model.User {
-	db := database.GetDB()
-	user := &model.User{}
-	err := db.Model(model.User{}).
-		Where("id = ?", id).
-		First(user).
-		Error
-	if err == gorm.ErrRecordNotFound {
-		return nil
+		if gotp.NewDefaultTOTP(twoFactorToken).Now() != twoFactorCode {
+			return nil
+		}
 	}
+
 	return user
 }
 
-func (s *UserService) CheckSecretExistence() (bool, error) {
+func (s *UserService) UpdateUser(id int, username string, password string) error {
 	db := database.GetDB()
+	hashedPassword, err := crypto.HashPasswordAsBcrypt(password)
 
-	var count int64
-	err := db.Model(model.User{}).
-		Where("login_secret IS NOT NULL").
-		Count(&count).
-		Error
 	if err != nil {
-		return false, err
+		return err
 	}
 
-	return count > 0, nil
+	return db.Model(model.User{}).
+		Where("id = ?", id).
+		Updates(map[string]any{"username": username, "password": hashedPassword}).
+		Error
 }
 
 func (s *UserService) UpdateFirstUser(username string, password string) error {