kra-new/internal/server/handler/public.go

package handler

import (
	"context"
	"errors"
	"strconv"
	"strings"
	"time"

	"kra/internal/biz"
	"kra/internal/service"
	"kra/internal/service/dto"

	"github.com/gin-gonic/gin"
	"github.com/mojocn/base64Captcha"
)

type Public struct {
	auth     *service.AuthService
	system   *service.SystemConfigService
	settings *service.SecurityService
	store    *captchaStore
}

func NewPublic(auth *service.AuthService, system *service.SystemConfigService, settings *service.SecurityService) *Public {
	return &Public{auth: auth, system: system, settings: settings, store: &captchaStore{service: settings}}
}

func (h *Public) captchaConfig() (int, int, int) {
	config := h.settings.CaptchaSettings()
	return config.KeyLong, config.ImageWidth, config.ImageHeight
}

func (h *Public) Captcha(c *gin.Context) {
	keyLong, width, height := h.captchaConfig()
	security, err := h.settings.CurrentSecurity(c.Request.Context())
	if err != nil || security == nil {
		Fail(c, "安全服务暂不可用")
		return
	}
	openCaptcha := true
	if security != nil {
		keyLong, width, height = security.KeyLong, security.ImgWidth, security.ImgHeight
		ttl := time.Duration(security.CaptchaTimeout) * time.Second
		failures, counterErr := h.settings.EnsureLoginIPCounter(c.Request.Context(), c.ClientIP(), ttl)
		if counterErr != nil {
			Fail(c, "安全服务暂不可用")
			return
		}
		openCaptcha = security.CaptchaOpen == 0 || failures > security.CaptchaOpen
	}
	driver := base64Captcha.NewDriverDigit(height, width, keyLong, 0.7, 80)
	id, picture, _, err := base64Captcha.NewCaptcha(driver, h.store).Generate()
	if err != nil {
		Fail(c, "验证码获取失败")
		return
	}
	Write(c, CodeSuccess, gin.H{"captchaId": id, "picPath": picture, "captchaLength": keyLong, "openCaptcha": openCaptcha}, "验证码获取成功")
}

func (h *Public) Login(c *gin.Context) {
	var req dto.LoginRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		Fail(c, err.Error())
		return
	}
	if req.Username == "" {
		Fail(c, "Username值不能为空")
		return
	}
	if req.Password == "" {
		Fail(c, "Password值不能为空")
		return
	}
	result, err := h.auth.Login(c.Request.Context(), &req, c.ClientIP(), c.Request.UserAgent())
	if err != nil {
		switch {
		case errors.Is(err, biz.ErrAccountLocked):
			var locked *biz.AccountLockedError
			errors.As(err, &locked)
			minutes := 0
			if locked != nil {
				minutes = locked.Minutes
			}
			Fail(c, "账号已锁定，请 "+strconv.Itoa(minutes)+" 分钟后再试")
		case errors.Is(err, biz.ErrCaptchaInvalid):
			Fail(c, "验证码错误")
		case errors.Is(err, biz.ErrUserDisabled):
			Fail(c, "用户被禁止登录")
		case errors.Is(err, biz.ErrTokenIssue):
			Fail(c, "获取token失败")
		case errors.Is(err, biz.ErrLoginState):
			Fail(c, "设置登录状态失败")
		case errors.Is(err, biz.ErrTokenRevoke):
			Fail(c, "jwt作废失败")
		default:
			Fail(c, "用户名不存在或者密码错误")
		}
		return
	}
	maxAge := int(time.UnixMilli(result.ExpiresAt).Unix() - time.Now().Unix())
	SetTokenCookie(c, result.Token, maxAge)
	Write(c, CodeSuccess, result, "登录成功")
}

func (h *Public) CheckDatabase(c *gin.Context) {
	initialized, err := h.system.IsInitialized(c.Request.Context())
	if err != nil {
		Fail(c, "数据库状态检查失败")
		return
	}
	message := "数据库无需初始化"
	if !initialized {
		message = "前往初始化数据库"
	}
	Write(c, CodeSuccess, gin.H{"needInit": !initialized}, message)
}

func (h *Public) InitializeDatabase(engine *gin.Engine) gin.HandlerFunc {
	return func(c *gin.Context) {
		initialized, err := h.system.IsInitialized(c.Request.Context())
		if err != nil {
			Fail(c, "数据库状态检查失败")
			return
		}
		if initialized {
			Fail(c, "已存在数据库配置")
			return
		}
		var input dto.DatabaseInitRequest
		if c.ShouldBindJSON(&input) != nil {
			Fail(c, "参数校验不通过")
			return
		}
		routes := engine.Routes()
		values := make([]dto.Route, 0, len(routes))
		for _, route := range routes {
			values = append(values, dto.Route{Path: route.Path, Method: route.Method})
		}
		if err := h.system.InitializeRoutes(c.Request.Context(), &input, values); err != nil {
			Fail(c, "自动创建数据库失败，请查看后台日志，检查后在进行初始化")
			return
		}
		Write(c, CodeSuccess, gin.H{}, "自动创建数据库成功")
	}
}

type captchaStore struct {
	service *service.SecurityService
}

func (s *captchaStore) Set(id, value string) error {
	expiration := s.service.CaptchaSettings().StoreExpiration
	return s.service.SetCaptcha(context.Background(), id, value, expiration)
}
func (s *captchaStore) Get(id string, clear bool) string {
	value, ok, err := s.service.GetCaptcha(context.Background(), id)
	if err != nil || !ok {
		return ""
	}
	if clear {
		_ = s.service.DeleteCaptcha(context.Background(), id)
	}
	return value
}
func (s *captchaStore) Verify(id, answer string, clear bool) bool {
	return id != "" && answer != "" && strings.EqualFold(s.Get(id, clear), answer)
}