user-system/internal/api/handler/totp_handler.go

package handler

import (
	"net/http"

	"github.com/gin-gonic/gin"

	"github.com/user-management-system/internal/service"
)

// TOTPHandler handles TOTP 2FA requests
type TOTPHandler struct {
	authService *service.AuthService
	totpService *service.TOTPService
}

// NewTOTPHandler creates a new TOTPHandler
func NewTOTPHandler(authService *service.AuthService, totpService *service.TOTPService) *TOTPHandler {
	return &TOTPHandler{
		authService: authService,
		totpService: totpService,
	}
}

// GetTOTPStatus 获取TOTP状态
// @Summary 获取TOTP状态
// @Description 获取当前用户的TOTP两步验证状态
// @Tags 两步验证
// @Produce json
// @Security BearerAuth
// @Success 200 {object} Response{data=TOTPStatusResponse} "TOTP状态"
// @Failure 401 {object} Response "未认证"
// @Router /api/v1/auth/totp/status [get]
func (h *TOTPHandler) GetTOTPStatus(c *gin.Context) {
	userID, ok := getUserIDFromContext(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"code": 401, "message": "unauthorized"})
		return
	}

	enabled, err := h.totpService.GetTOTPStatus(c.Request.Context(), userID)
	if err != nil {
		handleError(c, err)
		return
	}

	c.JSON(http.StatusOK, gin.H{"code": 0, "message": "success", "data": gin.H{"enabled": enabled}})
}

// SetupTOTP 设置 TOTP
// @Summary 设置 TOTP 两步验证
// @Description 为当前用户设置 TOTP 两步验证，返回密钥和二维码
// @Tags 两步验证
// @Accept json
// @Produce json
// @Security BearerAuth
// @Success 200 {object} Response{data=TOTPSetupResponse} "TOTP设置信息"
// @Failure 401 {object} Response "未认证"
// @Failure 500 {object} Response "服务器错误"
// @Router /api/v1/auth/totp/setup [post]
func (h *TOTPHandler) SetupTOTP(c *gin.Context) {
	userID, ok := getUserIDFromContext(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"code": 401, "message": "unauthorized"})
		return
	}

	resp, err := h.totpService.SetupTOTP(c.Request.Context(), userID)
	if err != nil {
		handleError(c, err)
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"code":    0,
		"message": "success",
		"data": gin.H{
			"secret":         resp.Secret,
			"qr_code_base64": resp.QRCodeBase64,
			"recovery_codes": resp.RecoveryCodes,
		},
	})
}

// EnableTOTP 启用 TOTP
// @Summary 启用 TOTP 两步验证
// @Description 输入验证码启用 TOTP 两步验证
// @Tags 两步验证
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param request body EnableTOTPRequest true "验证码"
// @Success 200 {object} Response "启用成功"
// @Failure 400 {object} Response "请求参数错误"
// @Failure 401 {object} Response "未认证或验证码错误"
// @Failure 500 {object} Response "服务器错误"
// @Router /api/v1/auth/totp/enable [post]
func (h *TOTPHandler) EnableTOTP(c *gin.Context) {
	userID, ok := getUserIDFromContext(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"code": 401, "message": "unauthorized"})
		return
	}

	var req struct {
		Code string `json:"code" binding:"required"`
	}

	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"code": 400, "message": err.Error()})
		return
	}

	if err := h.totpService.EnableTOTP(c.Request.Context(), userID, req.Code); err != nil {
		handleError(c, err)
		return
	}

	c.JSON(http.StatusOK, gin.H{"code": 0, "message": "success"})
}

// DisableTOTP 禁用 TOTP
// @Summary 禁用 TOTP 两步验证
// @Description 输入验证码禁用 TOTP 两步验证
// @Tags 两步验证
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param request body DisableTOTPRequest true "验证码"
// @Success 200 {object} Response "禁用成功"
// @Failure 400 {object} Response "请求参数错误"
// @Failure 401 {object} Response "未认证或验证码错误"
// @Failure 500 {object} Response "服务器错误"
// @Router /api/v1/auth/totp/disable [post]
func (h *TOTPHandler) DisableTOTP(c *gin.Context) {
	userID, ok := getUserIDFromContext(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"code": 401, "message": "unauthorized"})
		return
	}

	var req struct {
		Code string `json:"code" binding:"required"`
	}

	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"code": 400, "message": err.Error()})
		return
	}

	if err := h.totpService.DisableTOTP(c.Request.Context(), userID, req.Code); err != nil {
		handleError(c, err)
		return
	}

	c.JSON(http.StatusOK, gin.H{"code": 0, "message": "success"})
}

// VerifyTOTP 验证 TOTP
// @Summary 验证 TOTP 验证码
// @Description 在登录或其他敏感操作时验证 TOTP 验证码
// @Tags 两步验证
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param request body VerifyTOTPRequest true "验证码"
// @Success 200 {object} Response{data=VerifyTOTPResponse} "验证结果"
// @Failure 400 {object} Response "请求参数错误"
// @Failure 401 {object} Response "未认证或验证码错误"
// @Failure 500 {object} Response "服务器错误"
// @Router /api/v1/auth/totp/verify [post]
func (h *TOTPHandler) VerifyTOTP(c *gin.Context) {
	userID, ok := getUserIDFromContext(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"code": 401, "message": "unauthorized"})
		return
	}

	var req struct {
		Code     string `json:"code" binding:"required"`
		DeviceID string `json:"device_id,omitempty"`
	}

	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"code": 400, "message": err.Error()})
		return
	}

	if err := h.authService.VerifyTOTP(c.Request.Context(), userID, req.Code, req.DeviceID); err != nil {
		handleError(c, err)
		return
	}

	c.JSON(http.StatusOK, gin.H{"code": 0, "message": "success", "data": gin.H{"verified": true}})
}