From 896e242827c8a15b728ccb148aa4706d81e5768f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=80=97=E5=AD=90?= <haozi@loli.email>
Date: Wed, 14 May 2025 04:21:17 +0800
Subject: [PATCH] =?UTF-8?q?feat:=202fa=E5=91=BD=E4=BB=A4=E8=A1=8C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/data/user.go   |  2 +-
 internal/route/cli.go   |  5 +++++
 internal/service/cli.go | 46 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/internal/data/user.go b/internal/data/user.go
index 9aa62e0f..efa41f47 100644
--- a/internal/data/user.go
+++ b/internal/data/user.go
@@ -145,7 +145,7 @@ func (r *userRepo) UpdateTwoFA(id uint, code, secret string) error {
 	}
 
 	// 保存前先验证一次，防止错误开启
-	if !totp.Validate(code, secret) {
+	if secret != "" && !totp.Validate(code, secret) {
 		return errors.New(r.t.Get("invalid 2fa code"))
 	}
 
diff --git a/internal/route/cli.go b/internal/route/cli.go
index e0f5a32a..b7870f11 100644
--- a/internal/route/cli.go
+++ b/internal/route/cli.go
@@ -75,6 +75,11 @@ func (route *Cli) Commands() []*cli.Command {
 					Usage:  route.t.Get("Change user password"),
 					Action: route.cli.UserPassword,
 				},
+				{
+					Name:   "2fa",
+					Usage:  route.t.Get("Change user 2FA"),
+					Action: route.cli.UserTwoFA,
+				},
 			},
 		},
 		{
diff --git a/internal/service/cli.go b/internal/service/cli.go
index e08fed2b..4acc0b6c 100644
--- a/internal/service/cli.go
+++ b/internal/service/cli.go
@@ -1,10 +1,13 @@
 package service
 
 import (
+	"bufio"
 	"context"
 	"errors"
 	"fmt"
+	stdos "os"
 	"path/filepath"
+	"strings"
 	"time"
 
 	"github.com/go-rat/utils/collect"
@@ -260,6 +263,49 @@ func (s *CliService) UserPassword(ctx context.Context, cmd *cli.Command) error {
 	return nil
 }
 
+func (s *CliService) UserTwoFA(ctx context.Context, cmd *cli.Command) error {
+	user := new(biz.User)
+	username := cmd.Args().Get(0)
+	if username == "" {
+		return errors.New(s.t.Get("Username cannot be empty"))
+	}
+
+	if err := s.db.Where("username", username).First(user).Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return errors.New(s.t.Get("User not exists"))
+		} else {
+			return errors.New(s.t.Get("Failed to get user: %v", err))
+		}
+	}
+
+	// 已开启，关闭2FA
+	if user.TwoFA != "" {
+		user.TwoFA = ""
+		if err := s.db.Save(user).Error; err != nil {
+			return errors.New(s.t.Get("Failed to change 2FA status: %v", err))
+		}
+		fmt.Println(s.t.Get("2FA disabled for user %s", username))
+		return nil
+	}
+	// 未开启，开启2FA
+	_, url, secret, err := s.userRepo.GenerateTwoFA(user.ID)
+	if err != nil {
+		return errors.New(s.t.Get("Failed to generate 2FA: %v", err))
+	}
+	fmt.Println(s.t.Get("2FA url: %s", url))
+	reader := bufio.NewReader(stdos.Stdin)
+	fmt.Print(s.t.Get("Please enter the 2FA code: "))
+	code, err := reader.ReadString('\n')
+	if err != nil {
+		return errors.New(s.t.Get("Failed to read input: %v", err))
+	}
+	if err = s.userRepo.UpdateTwoFA(user.ID, strings.TrimSpace(code), secret); err != nil {
+		return errors.New(s.t.Get("Failed to update 2FA: %v", err))
+	}
+
+	return nil
+}
+
 func (s *CliService) HTTPSOn(ctx context.Context, cmd *cli.Command) error {
 	config := new(types.PanelConfig)
 	raw, err := io.Read("/usr/local/etc/panel/config.yml")