hfb_sys/backend/internal/modules/adminauth/repository.go

package adminauth

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"html"
	"math/big"
	"strings"
	"time"

	"hfb_sys/backend/internal/model"
	"hfb_sys/backend/internal/modules/auth"

	"github.com/redis/go-redis/v9"
	"golang.org/x/crypto/bcrypt"
	"gorm.io/gorm"
)

const (
	defaultAdminUsername = "admin"
	defaultAdminPassword = "admin123456"
	captchaTTL           = 3 * time.Minute
)

type Repository struct {
	db    *gorm.DB
	redis *redis.Client
	jwt   *auth.JWTManager
}

func NewRepository(db *gorm.DB, redis *redis.Client, jwt *auth.JWTManager) *Repository {
	return &Repository{db: db, redis: redis, jwt: jwt}
}

func (r *Repository) Captcha() (*CaptchaDTO, error) {
	if r.redis == nil {
		return nil, ErrDependencyUnavailable
	}
	captchaID, err := randomToken(16)
	if err != nil {
		return nil, err
	}
	code, err := randomCaptchaCode(4)
	if err != nil {
		return nil, err
	}
	ctx := context.Background()
	if err := r.redis.Set(ctx, captchaKey(captchaID), strings.ToUpper(code), captchaTTL).Err(); err != nil {
		return nil, err
	}
	return &CaptchaDTO{
		CaptchaID: captchaID,
		Image:     captchaImageDataURL(code),
		ExpiresIn: int64(captchaTTL.Seconds()),
	}, nil
}

func (r *Repository) Login(username string, password string, captchaID string, captchaCode string) (LoginResult, error) {
	if err := r.verifyCaptcha(captchaID, captchaCode); err != nil {
		return LoginResult{}, err
	}
	if err := r.ensureDefaultAdmin(); err != nil {
		return LoginResult{}, err
	}
	var admin model.AdminUser
	if err := r.db.Where("username = ?", username).First(&admin).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return LoginResult{}, ErrInvalidCredential
		}
		return LoginResult{}, err
	}
	if admin.Status != "active" {
		return LoginResult{}, ErrAdminDisabled
	}
	if err := bcrypt.CompareHashAndPassword([]byte(admin.PasswordHash), []byte(password)); err != nil {
		return LoginResult{}, ErrInvalidCredential
	}
	now := time.Now()
	admin.LastLoginAt = &now
	if err := r.db.Save(&admin).Error; err != nil {
		return LoginResult{}, err
	}
	tokens, err := r.jwt.GenerateSubjectPair(admin.ID, admin.Username, "admin")
	if err != nil {
		return LoginResult{}, err
	}
	dto := toDTO(admin)
	r.loadRolesAndPerms(&dto)
	return LoginResult{Admin: dto, Tokens: tokens}, nil
}

func (r *Repository) verifyCaptcha(captchaID string, captchaCode string) error {
	if r.redis == nil {
		return ErrDependencyUnavailable
	}
	ctx := context.Background()
	key := captchaKey(captchaID)
	stored, err := r.redis.Get(ctx, key).Result()
	if errors.Is(err, redis.Nil) {
		return ErrCaptchaInvalid
	}
	if err != nil {
		return err
	}
	_ = r.redis.Del(ctx, key).Err()
	if strings.ToUpper(strings.TrimSpace(captchaCode)) != stored {
		return ErrCaptchaInvalid
	}
	return nil
}

func (r *Repository) FindByID(id uint64) (*AdminDTO, error) {
	var admin model.AdminUser
	if err := r.db.First(&admin, id).Error; err != nil {
		return nil, err
	}
	if admin.Status != "active" {
		return nil, ErrAdminDisabled
	}
	dto := toDTO(admin)
	r.loadRolesAndPerms(&dto)
	return &dto, nil
}

func (r *Repository) UpdateSupportStatus(adminID uint64, status string) error {
	if r.db == nil {
		return ErrDependencyUnavailable
	}
	if status != "online" && status != "offline" && status != "busy" {
		return errors.New("invalid support status")
	}
	return r.db.Model(&model.AdminUser{}).
		Where("id = ?", adminID).
		Update("support_status", status).Error
}

func (r *Repository) ensureDefaultAdmin() error {
	var count int64
	if err := r.db.Model(&model.AdminUser{}).Count(&count).Error; err != nil {
		return err
	}
	if count > 0 {
		return nil
	}
	hash, err := bcrypt.GenerateFromPassword([]byte(defaultAdminPassword), bcrypt.DefaultCost)
	if err != nil {
		return err
	}
	admin := model.AdminUser{
		Username:     defaultAdminUsername,
		PasswordHash: string(hash),
		Nickname:     "超级管理员",
		Status:       "active",
	}
	if err := r.db.Create(&admin).Error; err != nil {
		return err
	}

	// 自动关联 super_admin 角色
	var superAdminRole model.Role
	if err := r.db.Where("code = ?", "super_admin").First(&superAdminRole).Error; err == nil {
		r.db.Create(&model.AdminUserRole{
			AdminID: admin.ID,
			RoleID:  superAdminRole.ID,
		})
	}

	return nil
}

func toDTO(admin model.AdminUser) AdminDTO {
	return AdminDTO{
		ID:            admin.ID,
		Username:      admin.Username,
		Nickname:      admin.Nickname,
		Status:        admin.Status,
		SupportStatus: admin.SupportStatus,
		LastLoginAt:   admin.LastLoginAt,
	}
}

func (r *Repository) loadRolesAndPerms(dto *AdminDTO) {
	if r.db == nil {
		return
	}
	// 加载角色
	var roles []RoleDTO
	r.db.Table("roles").
		Joins("JOIN admin_user_roles aur ON aur.role_id = roles.id").
		Where("aur.admin_user_id = ?", dto.ID).
		Find(&roles)
	dto.Roles = roles

	// 加载权限
	var permCodes []string
	r.db.Table("permissions").
		Select("DISTINCT permissions.code").
		Joins("JOIN role_permissions rp ON rp.permission_id = permissions.id").
		Joins("JOIN admin_user_roles aur ON aur.role_id = rp.role_id").
		Where("aur.admin_user_id = ?", dto.ID).
		Pluck("code", &permCodes)
	dto.Permissions = permCodes

	// 缓存权限到 Redis
	if r.redis != nil && len(permCodes) > 0 {
		ctx := context.Background()
		key := fmt.Sprintf("admin:perms:%d", dto.ID)
		raw, _ := json.Marshal(permCodes)
		r.redis.Set(ctx, key, string(raw), 2*time.Hour)
	}
}

func captchaKey(id string) string {
	return "admin:captcha:" + id
}

func randomToken(length int) (string, error) {
	buf := make([]byte, length)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(buf), nil
}

func randomCaptchaCode(length int) (string, error) {
	const alphabet = "23456789ABCDEFGHJKLMNPQRSTUVWXYZ"
	result := make([]byte, length)
	for i := range result {
		n, err := rand.Int(rand.Reader, big.NewInt(int64(len(alphabet))))
		if err != nil {
			return "", err
		}
		result[i] = alphabet[n.Int64()]
	}
	return string(result), nil
}

func captchaImageDataURL(code string) string {
	safeCode := html.EscapeString(strings.ToUpper(code))
	svg := fmt.Sprintf(`<svg xmlns="http://www.w3.org/2000/svg" width="132" height="44" viewBox="0 0 132 44">
<rect width="132" height="44" rx="8" fill="#eef5f7"/>
<path d="M8 32 C32 2, 62 52, 124 12" stroke="#0f766e" stroke-width="2" fill="none" opacity=".28"/>
<path d="M10 13 C42 42, 86 0, 122 31" stroke="#2563eb" stroke-width="2" fill="none" opacity=".22"/>
<text x="66" y="29" text-anchor="middle" font-family="Menlo,Consolas,monospace" font-size="24" font-weight="700" letter-spacing="4" fill="#111827">%s</text>
</svg>`, safeCode)
	return "data:image/svg+xml;base64," + base64.StdEncoding.EncodeToString([]byte(svg))
}