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

package handler

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

	"kra/internal/biz"
	"kra/internal/conf"
	"kra/internal/server/httpx"
	"kra/internal/service"
	"kra/internal/service/dto"
	"kra/internal/worker"

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

type Public struct {
	system    *service.SystemService
	settings  *service.SettingsService
	audit     *service.AuditService
	scheduler *worker.TaskScheduler
	runtime   *conf.Runtime
	store     *captchaStore
}

func NewPublic(runtime *conf.Runtime, system *service.SystemService, settings *service.SettingsService, audit *service.AuditService, scheduler *worker.TaskScheduler) *Public {
	return &Public{system: system, settings: settings, audit: audit, scheduler: scheduler, runtime: runtime, store: &captchaStore{service: system, runtime: runtime}}
}

func (h *Public) captchaConfig() (int, int, int) {
	keyLong, width, height := 6, 240, 80
	config := h.runtime.Admin()
	if config != nil && config.Captcha != nil {
		if config.Captcha.KeyLong > 0 {
			keyLong = int(config.Captcha.KeyLong)
		}
		if config.Captcha.ImgWidth > 0 {
			width = int(config.Captcha.ImgWidth)
		}
		if config.Captcha.ImgHeight > 0 {
			height = int(config.Captcha.ImgHeight)
		}
	}
	return keyLong, width, height
}

func (h *Public) Captcha(c *gin.Context) {
	keyLong, width, height := h.captchaConfig()
	security, _ := h.settings.CurrentSecurity(c.Request.Context())
	openCaptcha := true
	if security != nil {
		keyLong, width, height = security.KeyLong, security.ImgWidth, security.ImgHeight
		ttl := time.Duration(security.CaptchaTimeout) * time.Second
		failures, _ := ensureLoginIPCounter(c.Request.Context(), h.system, c.ClientIP(), ttl)
		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 {
		httpx.Fail(c, "验证码获取失败")
		return
	}
	httpx.Write(c, httpx.CodeSuccess, gin.H{"captchaId": id, "picPath": picture, "captchaLength": keyLong, "openCaptcha": openCaptcha}, "验证码获取成功")
}

func (h *Public) Login(c *gin.Context) {
	var req dto.LoginRequest
	if c.ShouldBindJSON(&req) != nil || req.Username == "" || req.Password == "" {
		httpx.Fail(c, "请输入用户名和密码")
		return
	}
	security, _ := h.settings.CurrentSecurity(c.Request.Context())
	if security != nil && security.LockEnable {
		if _, locked, _ := h.system.CacheGet(c.Request.Context(), "login:lock:"+req.Username); locked {
			httpx.Fail(c, "账号已锁定，请 "+strconv.Itoa(security.LockDuration)+" 分钟后再试")
			_ = h.audit.RecordLoginRequest(c.Request.Context(), &dto.LoginLogRequest{Username: req.Username, IP: c.ClientIP(), Status: false, ErrorMessage: "账号已锁定", Agent: c.Request.UserAgent()})
			return
		}
	}
	requireCaptcha := security == nil || security.CaptchaOpen == 0
	ipTTL := time.Hour
	if security != nil {
		if security.CaptchaTimeout > 0 {
			ipTTL = time.Duration(security.CaptchaTimeout) * time.Second
		}
		failures, _ := ensureLoginIPCounter(c.Request.Context(), h.system, c.ClientIP(), ipTTL)
		requireCaptcha = security.CaptchaOpen == 0 || failures > security.CaptchaOpen
	}
	if requireCaptcha && (req.CaptchaID == "" || req.Captcha == "" || !h.store.Verify(req.CaptchaID, req.Captcha, true)) {
		_, _ = h.system.CacheIncrement(c.Request.Context(), c.ClientIP(), ipTTL)
		_ = h.audit.RecordLoginRequest(c.Request.Context(), &dto.LoginLogRequest{Username: req.Username, IP: c.ClientIP(), Status: false, ErrorMessage: "验证码错误", Agent: c.Request.UserAgent()})
		httpx.Fail(c, "验证码错误")
		return
	}
	result, err := h.system.Login(c.Request.Context(), req.Username, req.Password)
	if err != nil {
		_, _ = h.system.CacheIncrement(c.Request.Context(), c.ClientIP(), ipTTL)
		if errors.Is(err, biz.ErrUserDisabled) {
			var disabled *service.UserDisabledError
			errors.As(err, &disabled)
			userID := uint(0)
			if disabled != nil {
				userID = disabled.UserID
			}
			_ = h.audit.RecordLoginRequest(c.Request.Context(), &dto.LoginLogRequest{Username: req.Username, IP: c.ClientIP(), Status: false, ErrorMessage: "用户被禁止登录", Agent: c.Request.UserAgent(), UserID: userID})
			httpx.Fail(c, "用户被禁止登录")
			return
		}
		_ = h.audit.RecordLoginRequest(c.Request.Context(), &dto.LoginLogRequest{Username: req.Username, IP: c.ClientIP(), Status: false, ErrorMessage: "用户名不存在或者密码错误", Agent: c.Request.UserAgent()})
		if security != nil && security.LockEnable {
			lockTTL := time.Duration(security.LockDuration) * time.Minute
			failures, _ := h.system.CacheIncrement(c.Request.Context(), "login:fail:"+req.Username, lockTTL)
			if int(failures) >= security.LockThreshold {
				_ = h.system.CacheSet(c.Request.Context(), "login:lock:"+req.Username, "1", lockTTL)
			}
		}
		httpx.Fail(c, "用户名不存在或者密码错误")
		return
	}
	userID, _ := result.User["ID"].(uint)
	_ = h.audit.RecordLoginRequest(c.Request.Context(), &dto.LoginLogRequest{Username: req.Username, IP: c.ClientIP(), Status: true, Agent: c.Request.UserAgent(), UserID: userID})
	_ = h.system.CacheDelete(c.Request.Context(), "login:fail:"+req.Username)
	_ = h.system.CacheDelete(c.Request.Context(), "login:lock:"+req.Username)
	maxAge := int(time.Until(time.UnixMilli(result.ExpiresAt)).Seconds())
	httpx.SetTokenCookie(c, result.Token, maxAge)
	httpx.Write(c, httpx.CodeSuccess, result, "登录成功")
}

func ensureLoginIPCounter(ctx context.Context, system *service.SystemService, ip string, expiration time.Duration) (int, error) {
	value, exists, err := system.CacheGet(ctx, ip)
	if err != nil {
		return 0, err
	}
	if exists {
		return strconv.Atoi(value)
	}
	if expiration <= 0 {
		expiration = time.Hour
	}
	if err = system.CacheSet(ctx, ip, "1", expiration); err != nil {
		return 0, err
	}
	return 1, nil
}

func (h *Public) CheckDatabase(c *gin.Context) {
	initialized, err := h.system.IsInitialized(c.Request.Context())
	if err != nil {
		httpx.Fail(c, "数据库状态检查失败")
		return
	}
	message := "数据库无需初始化"
	if !initialized {
		message = "前往初始化数据库"
	}
	httpx.Write(c, httpx.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 {
			httpx.Fail(c, "数据库状态检查失败")
			return
		}
		if initialized {
			httpx.Fail(c, "数据库已初始化，无需重复初始化")
			return
		}
		var input service.DatabaseInit
		if c.ShouldBindJSON(&input) != nil {
			httpx.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 {
			httpx.Fail(c, "自动创建数据库失败: "+err.Error())
			return
		}
		if err := h.scheduler.Reload(c.Request.Context()); err != nil {
			httpx.Fail(c, "数据库已初始化，但定时任务加载失败:"+err.Error())
			return
		}
		httpx.Write(c, httpx.CodeSuccess, gin.H{}, "自动创建数据库成功")
	}
}

type captchaStore struct {
	service *service.SystemService
	runtime *conf.Runtime
}

func (s *captchaStore) key(id string) string { return "captcha:" + id }
func (s *captchaStore) Set(id, value string) error {
	expiration := 3 * time.Minute
	config := s.runtime.Admin()
	if config != nil && config.Captcha != nil && config.Captcha.StoreExpiration != nil {
		expiration = config.Captcha.StoreExpiration.AsDuration()
	}
	return s.service.CacheSet(context.Background(), s.key(id), value, expiration)
}
func (s *captchaStore) Get(id string, clear bool) string {
	value, ok, err := s.service.CacheGet(context.Background(), s.key(id))
	if err != nil || !ok {
		return ""
	}
	if clear {
		_ = s.service.CacheDelete(context.Background(), s.key(id))
	}
	return value
}
func (s *captchaStore) Verify(id, answer string, clear bool) bool {
	return id != "" && answer != "" && strings.EqualFold(s.Get(id, clear), answer)
}