From cff97d746c82064f8c23f0d4fcc727ca6c1bfa2b Mon Sep 17 00:00:00 2001
From: bjdgyc <bjdgyc@163.com>
Date: Tue, 3 Jan 2023 14:17:53 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E9=82=AE=E4=BB=B6=E5=86=85?=
 =?UTF-8?q?=E5=AE=B9=E7=9A=84otp=E5=9B=BE=E7=89=87?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/admin/api_user.go | 67 ++++++++++++++++++++++++----------------
 server/dbdata/db.go      |  4 +++
 2 files changed, 45 insertions(+), 26 deletions(-)

diff --git a/server/admin/api_user.go b/server/admin/api_user.go
index a869da3..b93ea5d 100644
--- a/server/admin/api_user.go
+++ b/server/admin/api_user.go
@@ -133,33 +133,44 @@ func UserDel(w http.ResponseWriter, r *http.Request) {
 
 func UserOtpQr(w http.ResponseWriter, r *http.Request) {
 	_ = r.ParseForm()
-	b64 := r.FormValue("b64")
+	b64S := r.FormValue("b64")
 	idS := r.FormValue("id")
 	id, _ := strconv.Atoi(idS)
-	var user dbdata.User
-	err := dbdata.One("Id", id, &user)
+
+	var b64 bool
+	if b64S == "1" {
+		b64 = true
+	}
+	data, err := userOtpQr(id, b64)
 	if err != nil {
-		RespError(w, RespInternalErr, err)
-		return
+		base.Error(err)
+	}
+	io.WriteString(w, data)
+}
+
+func userOtpQr(uid int, b64 bool) (string, error) {
+	var user dbdata.User
+	err := dbdata.One("Id", uid, &user)
+	if err != nil {
+		return "", err
 	}
 
 	issuer := url.QueryEscape(base.Cfg.Issuer)
 	qrstr := fmt.Sprintf("otpauth://totp/%s:%s?issuer=%s&secret=%s", issuer, user.Email, issuer, user.OtpSecret)
 	qr, _ := qrcode.New(qrstr, qrcode.High)
 
-	if b64 == "1" {
-		data, _ := qr.PNG(300)
-		s := base64.StdEncoding.EncodeToString(data)
-		_, err = fmt.Fprint(w, s)
+	if b64 {
+		data, err := qr.PNG(300)
 		if err != nil {
-			base.Error(err)
+			return "", err
 		}
-		return
-	}
-	err = qr.Write(300, w)
-	if err != nil {
-		base.Error(err)
+		s := base64.StdEncoding.EncodeToString(data)
+		return s, nil
 	}
+
+	buf := bytes.NewBuffer(nil)
+	err = qr.Write(300, buf)
+	return buf.String(), err
 }
 
 // 在线用户
@@ -190,12 +201,13 @@ func UserReline(w http.ResponseWriter, r *http.Request) {
 }
 
 type userAccountMailData struct {
-	Issuer   string
-	LinkAddr string
-	Group    string
-	Username string
-	PinCode  string
-	OtpImg   string
+	Issuer       string
+	LinkAddr     string
+	Group        string
+	Username     string
+	PinCode      string
+	OtpImg       string
+	OtpImgBase64 string
 }
 
 func userAccountMail(user *dbdata.User) error {
@@ -236,12 +248,15 @@ func userAccountMail(user *dbdata.User) error {
 		return err
 	}
 
+	otpData, _ := userOtpQr(user.Id, true)
+
 	data := userAccountMailData{
-		LinkAddr: setting.LinkAddr,
-		Group:    strings.Join(user.Groups, ","),
-		Username: user.Username,
-		PinCode:  user.PinCode,
-		OtpImg:   fmt.Sprintf("https://%s/otp_qr?id=%d&jwt=%s", setting.LinkAddr, user.Id, tokenString),
+		LinkAddr:     setting.LinkAddr,
+		Group:        strings.Join(user.Groups, ","),
+		Username:     user.Username,
+		PinCode:      user.PinCode,
+		OtpImg:       fmt.Sprintf("https://%s/otp_qr?id=%d&jwt=%s", setting.LinkAddr, user.Id, tokenString),
+		OtpImgBase64: "data:image/png;base64," + otpData,
 	}
 	w := bytes.NewBufferString("")
 	t, _ := template.New("auth_complete").Parse(htmlBody)
diff --git a/server/dbdata/db.go b/server/dbdata/db.go
index 676080e..24c9fae 100644
--- a/server/dbdata/db.go
+++ b/server/dbdata/db.go
@@ -149,8 +149,12 @@ const accountMail = `<p>您好:</p>
     用户组: <b>{{.Group}}</b> <br/>
     用户名: <b>{{.Username}}</b> <br/>
     用户PIN码: <b>{{.PinCode}}</b> <br/>
+    <!-- 
     用户动态码(3天后失效):<br/>
     <img src="{{.OtpImg}}"/>
+    -->
+    用户动态码(请妥善保存):<br/>
+    <img src="{{.OtpImgBase64}}"/>
 </p>
 <div>
     使用说明: