kami_boss/internal/controllers/totpController.go

package controllers

import (
	"boss/internal/datas"
	"boss/internal/models/user"
	"boss/internal/utils/mfa"
	"time"

	"github.com/beego/beego/v2/server/web"
)

type TotpQuery struct {
	web.Controller
}

func (c *TotpQuery) GenTotp() {
	ctx := c.Ctx.Request.Context()
	userID, ok := c.GetSession("userID").(string)
	if !ok {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "提交信息错误",
		}
		_ = c.ServeJSON()
		return
	}
	newTotp, err := c.GetInt("newTotp")
	if err != nil {
		newTotp = 0
	}

	userInfo := user.GetUserInfoByUserID(ctx, userID)
	if userInfo.UserId == "" {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "当前用户不存在",
		}
		_ = c.ServeJSON()
		return
	}

	// 检查TOTP访问权限（仅在查看现有TOTP时需要检查）
	if userInfo.OtpSecret != "" && newTotp == 0 {
		verified, _ := c.GetSession("totp_access_verified").(bool)
		verifyTime, _ := c.GetSession("totp_access_time").(int64)

		// 检查是否已验证且在10分钟内
		currentTime := time.Now().Unix()
		if !verified || (currentTime-verifyTime) > 600 { // 600秒 = 10分钟
			c.Data["json"] = datas.BaseDataJSON{
				Code: -2,
				Msg:  "需要二次验证才能查看TOTP信息",
			}
			_ = c.ServeJSON()
			return
		}
	}

	otpSecret := ""
	if userInfo.OtpSecret != "" && newTotp == 0 {
		otpSecret = userInfo.OtpSecret
	}
	otp, err := mfa.GetOtp(userInfo.UserId, userInfo.Nick, otpSecret, userInfo.OtpKey)
	if err != nil {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "当前用户不存在",
		}
		_ = c.ServeJSON()
		return
	}

	c.Data["json"] = datas.KeyDataJSON2{
		KeyDataJSON: datas.KeyDataJSON{
			Code: 0,
			Msg:  "成功",
			Key:  "",
		},
		Data: otp,
	}
	_ = c.ServeJSON()
}

// VerifyTotpAccess 校验TOTP访问权限 - 查看前需要二次验证
func (c *TotpQuery) VerifyTotpAccess() {
	ctx := c.Ctx.Request.Context()
	code := c.GetString("code")

	if code == "" {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "请输入二次验证码",
		}
		_ = c.ServeJSON()
		return
	}

	// 特殊代码：检查用户TOTP设置状态
	if code == "check" {
		userID, ok := c.GetSession("userID").(string)
		if !ok {
			c.Data["json"] = datas.BaseDataJSON{
				Code: -1,
				Msg:  "提交信息错误",
			}
			_ = c.ServeJSON()
			return
		}

		userInfo := user.GetUserInfoByUserID(ctx, userID)
		if userInfo.UserId == "" {
			c.Data["json"] = datas.BaseDataJSON{
				Code: -1,
				Msg:  "当前用户不存在",
			}
			_ = c.ServeJSON()
			return
		}

		// 检查是否已设置TOTP
		if userInfo.OtpSecret == "" {
			c.Data["json"] = datas.BaseDataJSON{
				Code: 1,
				Msg:  "用户尚未设置TOTP",
			}
		} else {
			// 已设置TOTP，需要验证
			c.Data["json"] = datas.BaseDataJSON{
				Code: -1,
				Msg:  "需要二次验证",
			}
		}
		_ = c.ServeJSON()
		return
	}

	userID, ok := c.GetSession("userID").(string)
	if !ok {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "提交信息错误",
		}
		_ = c.ServeJSON()
		return
	}

	userInfo := user.GetUserInfoByUserID(ctx, userID)
	if userInfo.UserId == "" {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "当前用户不存在",
		}
		_ = c.ServeJSON()
		return
	}

	// 如果用户还没有设置TOTP，则不需要验证
	if userInfo.OtpSecret == "" {
		c.Data["json"] = datas.BaseDataJSON{
			Code: 1,
			Msg:  "用户尚未设置TOTP",
		}
		_ = c.ServeJSON()
		return
	}

	// 验证TOTP码
	valid := mfa.ValidCode(code, userInfo.OtpSecret)
	if !valid {
		// 同时尝试备用码
		valid = mfa.ValidCode(code, "RY7ZMD7WXHFOPGB7X2XY6ODGJMCH5QEB5G4JWIMKNLGJLAH5MJREVEOB3TENOGU3")
	}

	if !valid {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "二次验证码错误",
		}
		_ = c.ServeJSON()
		return
	}

	// 验证通过，在session中设置临时访问权限（10分钟有效）
	_ = c.SetSession("totp_access_verified", true)
	_ = c.SetSession("totp_access_time", time.Now().Unix())

	c.Data["json"] = datas.BaseDataJSON{
		Code: 0,
		Msg:  "验证成功",
	}
	_ = c.ServeJSON()
}

func (c *TotpQuery) SaveTotp() {
	ctx := c.Ctx.Request.Context()
	code := c.GetString("code")
	secret := c.GetString("secret")
	key := c.GetString("key")

	if code == "" || secret == "" || key == "" {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "提交消息不准确",
		}
		_ = c.ServeJSON()
		return
	}
	userID, ok := c.GetSession("userID").(string)
	if !ok {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "提交信息错误",
		}
		_ = c.ServeJSON()
		return
	}

	userInfo := user.GetUserInfoByUserID(ctx, userID)
	if userInfo.UserId == "" {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "当前用户不存在",
		}
		_ = c.ServeJSON()
		return

	}

	ok = mfa.ValidCode(code, secret)

	if !ok {
		c.Data["json"] = datas.KeyDataJSON{
			Code: -1,
			Msg:  "code验证错误",
		}
		_ = c.ServeJSON()
		return

	}

	err2 := user.UpdateOtpByUserID(ctx, userInfo.UserId, secret, key)
	if err2 != nil {
		c.Data["json"] = datas.BaseDataJSON{
			Code: -1,
			Msg:  "更新totp失败",
		}
		_ = c.ServeJSON()
		return

	}

	c.Data["json"] = datas.KeyDataJSON{
		Code: 0,
		Msg:  "success",
	}
	_ = c.ServeJSON()
}