nl-pms-api/internal/service/admin_security.go

package service

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"net/url"
	"strings"
	"time"

	"github.com/pquerna/otp/totp"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"nl-pms-api/internal/commonservice"
	"nl-pms-api/internal/model"
)

const (
	adminStepupTTL    = 2 * time.Hour
	totpIssuer        = "年糕崽崽PMS"
	totpPendingPrefix = "pending:"
)

// AdminSecurityService 管理员 TOTP 与敏感操作 stepup。
type AdminSecurityService struct {
	DB *gorm.DB
}

type TOTPStatus struct {
	Enabled bool   `json:"enabled"`
	Pending bool   `json:"pending"`
	OTPAuth string `json:"otpauth,omitempty"`
	Secret  string `json:"secret,omitempty"`
}

// Status 返回 TOTP 绑定状态；若未启用且无 pending，不自动生成。
func (s *AdminSecurityService) Status(userID int64) (*TOTPStatus, error) {
	if userID != commonservice.AdminUserID {
		return nil, commonservice.Forbidden("FORBIDDEN")
	}
	var u model.User
	if err := s.DB.First(&u, userID).Error; err != nil {
		return nil, commonservice.Internal("QUERY_FAILED")
	}
	st := &TOTPStatus{Enabled: u.TOTPEnabled != 0}
	if !st.Enabled && strings.HasPrefix(u.TOTPSecret, totpPendingPrefix) {
		secret := strings.TrimPrefix(u.TOTPSecret, totpPendingPrefix)
		st.Pending = true
		st.Secret = secret
		st.OTPAuth = buildOTPAuth(u.Username, secret)
	}
	return st, nil
}

// SetupBegin 生成待确认的 TOTP 密钥（覆盖未确认的 pending）。
func (s *AdminSecurityService) SetupBegin(userID int64) (*TOTPStatus, error) {
	if userID != commonservice.AdminUserID {
		return nil, commonservice.Forbidden("FORBIDDEN")
	}
	var u model.User
	if err := s.DB.First(&u, userID).Error; err != nil {
		return nil, commonservice.Internal("QUERY_FAILED")
	}
	if u.TOTPEnabled != 0 {
		return nil, commonservice.Conflict("TOTP_ALREADY_ENABLED")
	}
	key, err := totp.Generate(totp.GenerateOpts{
		Issuer:      totpIssuer,
		AccountName: u.Username,
	})
	if err != nil {
		return nil, commonservice.Internal("INTERNAL")
	}
	secret := key.Secret()
	if err := s.DB.Model(&u).Update("totp_secret", totpPendingPrefix+secret).Error; err != nil {
		return nil, commonservice.Internal("SAVE_FAILED")
	}
	return &TOTPStatus{
		Enabled: false,
		Pending: true,
		Secret:  secret,
		OTPAuth: key.URL(),
	}, nil
}

// SetupConfirm 用动态码确认绑定。
func (s *AdminSecurityService) SetupConfirm(userID int64, code string) error {
	if userID != commonservice.AdminUserID {
		return commonservice.Forbidden("FORBIDDEN")
	}
	code = strings.TrimSpace(code)
	var u model.User
	if err := s.DB.First(&u, userID).Error; err != nil {
		return commonservice.Internal("QUERY_FAILED")
	}
	if u.TOTPEnabled != 0 {
		return commonservice.Conflict("TOTP_ALREADY_ENABLED")
	}
	if !strings.HasPrefix(u.TOTPSecret, totpPendingPrefix) {
		return commonservice.BadRequest("TOTP_SETUP_REQUIRED")
	}
	secret := strings.TrimPrefix(u.TOTPSecret, totpPendingPrefix)
	if !totp.Validate(code, secret) {
		return commonservice.BadRequest("TOTP_INVALID")
	}
	return s.DB.Model(&u).Updates(map[string]any{
		"totp_secret":  secret,
		"totp_enabled": 1,
	}).Error
}

// StepUp 校验动态码并签发 2h stepup token（绑定 IP）。
func (s *AdminSecurityService) StepUp(userID int64, code, clientIP string) (string, time.Time, error) {
	if userID != commonservice.AdminUserID {
		return "", time.Time{}, commonservice.Forbidden("FORBIDDEN")
	}
	code = strings.TrimSpace(code)
	var u model.User
	if err := s.DB.First(&u, userID).Error; err != nil {
		return "", time.Time{}, commonservice.Internal("QUERY_FAILED")
	}
	if u.TOTPEnabled == 0 || u.TOTPSecret == "" || strings.HasPrefix(u.TOTPSecret, totpPendingPrefix) {
		return "", time.Time{}, commonservice.BadRequest("TOTP_NOT_ENABLED")
	}
	if !totp.Validate(code, u.TOTPSecret) {
		return "", time.Time{}, commonservice.BadRequest("TOTP_INVALID")
	}
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", time.Time{}, commonservice.Internal("INTERNAL")
	}
	token := hex.EncodeToString(raw)
	sum := sha256.Sum256([]byte(token))
	expires := time.Now().UTC().Add(adminStepupTTL)
	row := model.AdminStepup{
		UserID:    userID,
		TokenHash: hex.EncodeToString(sum[:]),
		ClientIP:  strings.TrimSpace(clientIP),
		ExpiresAt: expires.Format(time.RFC3339),
	}
	if err := s.DB.Clauses(clause.OnConflict{
		Columns:   []clause.Column{{Name: "user_id"}},
		DoUpdates: clause.AssignmentColumns([]string{"token_hash", "client_ip", "expires_at"}),
	}).Create(&row).Error; err != nil {
		return "", time.Time{}, commonservice.Internal("SAVE_FAILED")
	}
	// 更新管理员最近登录 IP，便于审计
	_ = s.DB.Model(&u).Updates(map[string]any{
		"last_login_ip": row.ClientIP,
		"last_seen_at":  commonservice.NowRFC(),
	}).Error
	return token, expires, nil
}

// ClearStepUp 作废二次验证。
func (s *AdminSecurityService) ClearStepUp(userID int64) error {
	if userID != commonservice.AdminUserID {
		return commonservice.Forbidden("FORBIDDEN")
	}
	return s.DB.Delete(&model.AdminStepup{}, userID).Error
}

func buildOTPAuth(account, secret string) string {
	// otpauth URL：issuer/account 需编码
	label := url.PathEscape(totpIssuer) + ":" + url.PathEscape(account)
	q := url.Values{}
	q.Set("secret", secret)
	q.Set("issuer", totpIssuer)
	q.Set("algorithm", "SHA1")
	q.Set("digits", "6")
	q.Set("period", "30")
	return "otpauth://totp/" + label + "?" + q.Encode()
}