TOTP feature

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))
 	}

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)
+}

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))
 	}

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 {

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

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=

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
+}