secret token thanks to @HarlyquinForest

2 files changed, 47 insertions, 3 deletions

diff --git a/web/service/setting.go b/web/service/setting.go
index 9d92685e..6e305536 100644
--- a/web/service/setting.go
+++ b/web/service/setting.go
@@ -38,6 +38,7 @@ var defaultValueMap = map[string]string{
 	"tgRunTime":          "@daily",
 	"tgBotBackup":        "false",
 	"tgCpu":              "0",
+	"secretEnable":       "false",
 }
 
 type SettingService struct {
@@ -129,7 +130,13 @@ func (s *SettingService) GetAllSetting() (*entity.AllSetting, error) {
 
 func (s *SettingService) ResetSettings() error {
 	db := database.GetDB()
-	return db.Where("1 = 1").Delete(model.Setting{}).Error
+	err := db.Where("1 = 1").Delete(model.Setting{}).Error
+	if err != nil {
+		return err
+	}
+	return db.Model(model.User{}).
+		Where("1 = 1").
+		Update("login_secret", "").Error
 }
 
 func (s *SettingService) getSetting(key string) (*model.Setting, error) {
@@ -288,6 +295,14 @@ func (s *SettingService) SetgetTrafficDiff(value int) error {
 	return s.setInt("trafficDiff", value)
 }
 
+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 6da8bd1e..f1868424 100644
--- a/web/service/user.go
+++ b/web/service/user.go
@@ -25,12 +25,12 @@ func (s *UserService) GetFirstUser() (*model.User, error) {
 	return user, nil
 }
 
-func (s *UserService) CheckUser(username string, password string) *model.User {
+func (s *UserService) CheckUser(username string, password string, secret string) *model.User {
 	db := database.GetDB()
 
 	user := &model.User{}
 	err := db.Model(model.User{}).
-		Where("username = ? and password = ?", username, password).
+		Where("username = ? and password = ? and login_secret = ?", username, password, secret).
 		First(user).
 		Error
 	if err == gorm.ErrRecordNotFound {
@@ -50,6 +50,35 @@ func (s *UserService) UpdateUser(id int, username string, password string) error
 		Error
 }
 
+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
+}
+
+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
+	}
+	return user
+}
+
 func (s *UserService) UpdateFirstUser(username string, password string) error {
 	if username == "" {
 		return errors.New("username can not be empty")