✨ feat: totp close #9

pkg/api/account.go

@@ -1,16 +1,26 @@
 package api
 
 import (
-	"github.com/labstack/echo/v4"
+	"log"
+	"time"
+
 	"next-terminal/pkg/global"
 	"next-terminal/pkg/model"
+	"next-terminal/pkg/totp"
 	"next-terminal/pkg/utils"
-	"time"
+
+	"github.com/labstack/echo/v4"
 )
 
 type LoginAccount struct {
 	Username string `json:"username"`
 	Password string `json:"password"`
+	TOTP     string `json:"totp"`
+}
+
+type ConfirmTOTP struct {
+	Secret string `json:"secret"`
+	TOTP   string `json:"totp"`
 }
 
 type ChangePassword struct {
@@ -28,10 +38,17 @@ func LoginEndpoint(c echo.Context) error {
 	if err != nil {
 		return Fail(c, -1, "您输入的账号或密码不正确")
 	}
+
 	if err := utils.Encoder.Match([]byte(user.Password), []byte(loginAccount.Password)); err != nil {
 		return Fail(c, -1, "您输入的账号或密码不正确")
 	}
 
+	log.Println(user, loginAccount)
+
+	if !totp.Validate(loginAccount.TOTP, user.TOTPSecret) {
+		return Fail(c, -2, "您的TOTP不匹配")
+	}
+
 	token := utils.UUID()
 
 	global.Cache.Set(token, user, time.Minute*time.Duration(30))
@@ -47,6 +64,54 @@ func LogoutEndpoint(c echo.Context) error {
 	return Success(c, nil)
 }
 
+func ConfirmTOTPEndpoint(c echo.Context) error {
+	account, _ := GetCurrentAccount(c)
+
+	var confirmTOTP ConfirmTOTP
+	if err := c.Bind(&confirmTOTP); err != nil {
+		return err
+	}
+
+	if !totp.Validate(confirmTOTP.TOTP, confirmTOTP.Secret) {
+		return Fail(c, -1, "TOTP 验证失败，请重试")
+	}
+
+	u := &model.User{
+		TOTPSecret: confirmTOTP.Secret,
+	}
+
+	model.UpdateUserById(u, account.ID)
+
+	return Success(c, nil)
+}
+
+func ResetTOTPEndpoint(c echo.Context) error {
+	account, _ := GetCurrentAccount(c)
+
+	key, err := totp.NewTOTP(totp.GenerateOpts{
+		Issuer:      c.Request().Host,
+		AccountName: account.Username,
+	})
+	if err != nil {
+		return Fail(c, -1, err.Error())
+	}
+
+	qrcode, err := key.Image(200, 200)
+	if err != nil {
+		return Fail(c, -1, err.Error())
+	}
+
+	qrEncode, err := utils.ImageToBase64Encode(qrcode)
+	if err != nil {
+		return Fail(c, -1, err.Error())
+	}
+
+	return Success(c, map[string]string{
+		"qr":     qrEncode,
+		"secret": key.Secret(),
+	})
+}
+
 func ChangePasswordEndpoint(c echo.Context) error {
 	account, _ := GetCurrentAccount(c)
 

pkg/api/routes.go

@@ -1,11 +1,12 @@
 package api
 
 import (
-	"github.com/labstack/echo/v4"
-	"github.com/labstack/echo/v4/middleware"
 	"net/http"
 	"next-terminal/pkg/global"
 	"next-terminal/pkg/model"
+
+	"github.com/labstack/echo/v4"
+	"github.com/labstack/echo/v4/middleware"
 )
 
 const Token = "X-Auth-Token"
@@ -35,6 +36,8 @@ func SetupRoutes() *echo.Echo {
 
 	e.POST("/logout", LogoutEndpoint)
 	e.POST("/change-password", ChangePasswordEndpoint)
+	e.POST("/reset-totp", ResetTOTPEndpoint)
+	e.POST("/confirm-totp", ConfirmTOTPEndpoint)
 	e.GET("/info", InfoEndpoint)
 
 	users := e.Group("/users")

pkg/model/user.go

@@ -7,13 +7,14 @@ import (
 )
 
 type User struct {
-	ID       string         `gorm:"primary_key" json:"id"`
-	Username string         `json:"username"`
-	Password string         `json:"password"`
-	Nickname string         `json:"nickname"`
-	Online   bool           `json:"online"`
-	Enabled  bool           `json:"enabled"`
-	Created  utils.JsonTime `json:"created"`
+	ID         string         `gorm:"primary_key" json:"id"`
+	Username   string         `json:"username"`
+	Password   string         `json:"password"`
+	Nickname   string         `json:"nickname"`
+	TOTPSecret string         `json:"-"`
+	Online     bool           `json:"online"`
+	Enabled    bool           `json:"enabled"`
+	Created    utils.JsonTime `json:"created"`
 }
 
 func (r *User) TableName() string {
@@ -32,7 +33,6 @@ func FindAllUser() (o []User) {
 }
 
 func FindPageUser(pageIndex, pageSize int, username, nickname string) (o []User, total int64, err error) {
-
 	db := global.DB
 	if len(username) > 0 {
 		db = db.Where("username like ?", "%"+username+"%")

pkg/totp/totp.go

@@ -0,0 +1,19 @@
+package totp
+
+import (
+	otp_t "github.com/pquerna/otp"
+	totp_t "github.com/pquerna/otp/totp"
+)
+
+type GenerateOpts totp_t.GenerateOpts
+
+func NewTOTP(opt GenerateOpts) (*otp_t.Key, error) {
+	return totp_t.Generate(totp_t.GenerateOpts(opt))
+}
+
+func Validate(code string, secret string) bool {
+	if secret == "" {
+		return true
+	}
+	return totp_t.Validate(code, secret)
+}

pkg/utils/utils.go

@@ -1,14 +1,19 @@
 package utils
 
 import (
+	"bytes"
 	"database/sql/driver"
+	"encoding/base64"
 	"fmt"
-	"github.com/gofrs/uuid"
-	"golang.org/x/crypto/bcrypt"
+	"image"
+	"image/png"
 	"net"
 	"os"
 	"strconv"
 	"time"
+
+	"github.com/gofrs/uuid"
+	"golang.org/x/crypto/bcrypt"
 )
 
 type JsonTime struct {
@@ -81,6 +86,14 @@ func Tcping(ip string, port int) bool {
 	return true
 }
 
+func ImageToBase64Encode(img image.Image) (string, error) {
+	var buf bytes.Buffer
+	if err := png.Encode(&buf, img); err != nil {
+		return "", err
+	}
+	return base64.StdEncoding.EncodeToString(buf.Bytes()), nil
+}
+
 // 判断所给路径文件/文件夹是否存在
 func FileExists(path string) bool {
 	_, err := os.Stat(path) //os.Stat获取文件信息