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}})
}

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,
		},
	})
}

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"})
}

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"})
}

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}})
}