From d64a376496f3005495140176a0784d1b044afe4c Mon Sep 17 00:00:00 2001 From: Andrew Date: Sat, 1 Nov 2025 09:24:51 +0300 Subject: [PATCH] TOTP feature added --- api/apiHandler.go | 6 +++++ api/apiService.go | 50 ++++++++++++++++++++++++++++++++++++++++- api/apiV2Handler.go | 6 +++++ database/model/model.go | 10 +++++---- go.mod | 3 +++ go.sum | 4 ++++ service/user.go | 35 ++++++++++++++++++++++++----- 7 files changed, 104 insertions(+), 10 deletions(-) diff --git a/api/apiHandler.go b/api/apiHandler.go index 7e38a28f..edd341f7 100644 --- a/api/apiHandler.go +++ b/api/apiHandler.go @@ -56,6 +56,10 @@ func (a *APIHandler) postHandler(c *gin.Context) { case "deleteToken": a.ApiService.DeleteToken(c) a.apiv2.ReloadTokens() + case "enable2fa": + a.ApiService.Enable2FA(c) + case "disable2fa": + a.ApiService.Disable2FA(c) default: jsonMsg(c, "failed", common.NewError("unknown action: ", action)) } @@ -95,6 +99,8 @@ func (a *APIHandler) getHandler(c *gin.Context) { a.ApiService.GetDb(c) case "tokens": a.ApiService.GetTokens(c) + case "prepare2fa": + a.ApiService.Prepare2FA(c) default: jsonMsg(c, "failed", common.NewError("unknown action: ", action)) } diff --git a/api/apiService.go b/api/apiService.go index 2a8a8853..ec7f20cf 100644 --- a/api/apiService.go +++ b/api/apiService.go @@ -2,6 +2,7 @@ package api import ( "encoding/json" + "net/url" "strconv" "time" @@ -11,6 +12,8 @@ import ( "github.com/alireza0/s-ui/util" "github.com/gin-gonic/gin" + + "github.com/pquerna/otp/totp" ) type ApiService struct { @@ -261,7 +264,7 @@ func (a *ApiService) postActions(c *gin.Context) (string, json.RawMessage, error func (a *ApiService) Login(c *gin.Context) { remoteIP := getRemoteIp(c) - loginUser, err := a.UserService.Login(c.Request.FormValue("user"), c.Request.FormValue("pass"), remoteIP) + loginUser, err := a.UserService.Login(c.Request.FormValue("user"), c.Request.FormValue("pass"), c.Request.FormValue("passcode"), remoteIP) if err != nil { jsonMsg(c, "", err) return @@ -378,3 +381,48 @@ func (a *ApiService) DeleteToken(c *gin.Context) { err := a.UserService.DeleteToken(tokenId) jsonMsg(c, "", err) } + +func (a *ApiService) Prepare2FA(c *gin.Context) { + key, err := totp.Generate(totp.GenerateOpts{ + Issuer: "S-UI Panel", + AccountName: GetLoginUser(c), + }) + if err != nil { + jsonMsg(c, "", err) + return + } + secretUrl := key.URL() + jsonObj(c, secretUrl, nil) +} + +func (a *ApiService) Enable2FA(c *gin.Context) { + loginUser := GetLoginUser(c) + passcode := c.Request.FormValue("passcode") + secretUrl, err := url.Parse(c.Request.FormValue("secretUrl")) + if err != nil { + jsonMsg(c, "", err) + return + } + params, err := url.ParseQuery(secretUrl.RawQuery) + if err != nil { + jsonMsg(c, "", err) + return + } + secret := params.Get("secret") + err = a.UserService.Set2FAState(loginUser, secret, passcode, true) + if err != nil { + jsonMsg(c, "", err) + return + } + jsonMsg(c, "Enable2FA", err) +} + +func (a *ApiService) Disable2FA(c *gin.Context) { + loginUser := GetLoginUser(c) + err := a.UserService.Set2FAState(loginUser, "", "", false) + if err != nil { + jsonMsg(c, "", err) + return + } + jsonMsg(c, "Disable2FA", err) +} diff --git a/api/apiV2Handler.go b/api/apiV2Handler.go index 39cfd698..cdfe3835 100644 --- a/api/apiV2Handler.go +++ b/api/apiV2Handler.go @@ -51,6 +51,10 @@ func (a *APIv2Handler) postHandler(c *gin.Context) { a.ApiService.LinkConvert(c) case "importdb": a.ApiService.ImportDb(c) + case "enable2fa": + a.ApiService.Enable2FA(c) + case "disable2fa": + a.ApiService.Disable2FA(c) default: jsonMsg(c, "failed", common.NewError("unknown action: ", action)) } @@ -86,6 +90,8 @@ func (a *APIv2Handler) getHandler(c *gin.Context) { a.ApiService.GetKeypairs(c) case "getdb": a.ApiService.GetDb(c) + case "prepare2fa": + a.ApiService.Prepare2FA(c) default: jsonMsg(c, "failed", common.NewError("unknown action: ", action)) } diff --git a/database/model/model.go b/database/model/model.go index 56854601..3b7b5c7d 100644 --- a/database/model/model.go +++ b/database/model/model.go @@ -16,10 +16,12 @@ type Tls struct { } type User struct { - Id uint `json:"id" form:"id" gorm:"primaryKey;autoIncrement"` - Username string `json:"username" form:"username"` - Password string `json:"password" form:"password"` - LastLogins string `json:"lastLogin"` + Id uint `json:"id" form:"id" gorm:"primaryKey;autoIncrement"` + Username string `json:"username" form:"username"` + Password string `json:"password" form:"password"` + LastLogins string `json:"lastLogin"` + TOTPEnabled bool `json:"totp" gorm:"default:false"` + TOTPSecret string `json:"totp_secret" gorm:"size:64"` } type Client struct { diff --git a/go.mod b/go.mod index 5fa315e8..17f49a13 100644 --- a/go.mod +++ b/go.mod @@ -19,6 +19,8 @@ require ( gorm.io/gorm v1.31.0 ) +require github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect + require ( filippo.io/edwards25519 v1.1.0 // indirect github.com/ajg/form v1.5.1 // indirect @@ -97,6 +99,7 @@ require ( github.com/pelletier/go-toml/v2 v2.2.4 // indirect github.com/pierrec/lz4/v4 v4.1.21 // indirect github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect + github.com/pquerna/otp v1.4.0 github.com/prometheus-community/pro-bing v0.4.0 // indirect github.com/quic-go/qpack v0.5.1 // indirect github.com/quic-go/quic-go v0.54.0 // indirect diff --git a/go.sum b/go.sum index 82e4906d..8a281fb4 100644 --- a/go.sum +++ b/go.sum @@ -12,6 +12,8 @@ github.com/anytls/sing-anytls v0.0.8 h1:1u/fnH1HoeeMV5mX7/eUOjLBvPdkd1UJRmXiRi6V github.com/anytls/sing-anytls v0.0.8/go.mod h1:7rjN6IukwysmdusYsrV51Fgu1uW6vsrdd6ctjnEAln8= github.com/bits-and-blooms/bitset v1.13.0 h1:bAQ9OPNFYbGHV6Nez0tmNI0RiEu7/hxlYJRUA0wFAVE= github.com/bits-and-blooms/bitset v1.13.0/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8= +github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc h1:biVzkmvwrH8WK8raXaxBx6fRVTlJILwEwQGL1I/ByEI= +github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8= github.com/bytedance/sonic v1.14.0 h1:/OfKt8HFw0kh2rj8N0F6C/qPGRESq0BbaNZgcNXXzQQ= github.com/bytedance/sonic v1.14.0/go.mod h1:WoEbx8WTcFJfzCe0hbmyTGrfjt8PzNEBdxlNUO24NhA= github.com/bytedance/sonic/loader v0.3.0 h1:dskwH8edlzNMctoruo8FPTJDF3vLtDT0sXZwvZJyqeA= @@ -189,6 +191,8 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRI github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 h1:o4JXh1EVt9k/+g42oCprj/FisM4qX9L3sZB3upGN2ZU= github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= +github.com/pquerna/otp v1.4.0 h1:wZvl1TIVxKRThZIBiwOOHOGP/1+nZyWBil9Y2XNEDzg= +github.com/pquerna/otp v1.4.0/go.mod h1:dkJfzwRKNiegxyNb54X/3fLwhCynbMspSyWKnvi1AEg= github.com/prometheus-community/pro-bing v0.4.0 h1:YMbv+i08gQz97OZZBwLyvmmQEEzyfyrrjEaAchdy3R4= github.com/prometheus-community/pro-bing v0.4.0/go.mod h1:b7wRYZtCcPmt4Sz319BykUU241rWLe1VFXyiyWK/dH4= github.com/quic-go/qpack v0.5.1 h1:giqksBPnT/HDtZ6VhtFKgoLOWmlyo9Ei6u9PqzIMbhI= diff --git a/service/user.go b/service/user.go index aea2861f..680cacf5 100644 --- a/service/user.go +++ b/service/user.go @@ -2,12 +2,14 @@ package service import ( "encoding/json" + "errors" "time" "github.com/alireza0/s-ui/database" "github.com/alireza0/s-ui/database/model" "github.com/alireza0/s-ui/logger" "github.com/alireza0/s-ui/util/common" + "github.com/pquerna/otp/totp" ) type UserService struct { @@ -47,15 +49,15 @@ func (s *UserService) UpdateFirstUser(username string, password string) error { return db.Save(user).Error } -func (s *UserService) Login(username string, password string, remoteIP string) (string, error) { - user := s.CheckUser(username, password, remoteIP) +func (s *UserService) Login(username string, password string, passcode string, remoteIP string) (string, error) { + user := s.CheckUser(username, password, passcode, remoteIP) if user == nil { - return "", common.NewError("wrong user or password! IP: ", remoteIP) + return "", common.NewError("wrong user, password or passcode! IP: ", remoteIP) } return user.Username, nil } -func (s *UserService) CheckUser(username string, password string, remoteIP string) *model.User { +func (s *UserService) CheckUser(username string, password string, passcode string, remoteIP string) *model.User { db := database.GetDB() user := &model.User{} @@ -70,6 +72,11 @@ func (s *UserService) CheckUser(username string, password string, remoteIP strin return nil } + if user.TOTPEnabled && !totp.Validate(passcode, user.TOTPSecret) { + logger.Warning("check err:", errors.New("Invalid 2FA code"), " IP: ", remoteIP) + return nil + } + lastLoginTxt := time.Now().Format("2006-01-02 15:04:05") + " " + remoteIP err = db.Model(model.User{}). Where("username = ?", username). @@ -83,7 +90,7 @@ func (s *UserService) CheckUser(username string, password string, remoteIP strin func (s *UserService) GetUsers() (*[]model.User, error) { var users []model.User db := database.GetDB() - err := db.Model(model.User{}).Select("id,username,last_logins").Scan(&users).Error + err := db.Model(model.User{}).Select("id,username,last_logins,totp_enabled").Scan(&users).Error if err != nil { return nil, err } @@ -159,3 +166,21 @@ func (s *UserService) DeleteToken(id string) error { db := database.GetDB() return db.Model(model.Tokens{}).Where("id = ?", id).Delete(&model.Tokens{}).Error } + +func (s *UserService) Set2FAState(username string, secret string, passcode string, enabled bool) error { + db := database.GetDB() + user := &model.User{} + err := db.Model(model.User{}).Where("username = ?", username).First(user).Error + if err != nil || database.IsNotFound(err) { + return err + } + + if enabled && !totp.Validate(passcode, secret) { + return errors.New("Invalid passcode") + } + + user.TOTPEnabled = enabled + user.TOTPSecret = secret + + return db.Save(user).Error +}