GovAI/server/internal/service/twofa/service.go

// Package twofa 提供两步验证（2FA / TOTP）的业务编排（service 层）。
//
// 业务规则（生成密钥/备份码、校验、启用/关闭判定）集中在此，DB 操作通过 Store 接口
// 注入，便于单测。TOTP/备份码算法复用 pkg/auth。HTTP handler 仅做参数解析与响应。
package twofa

import (
	"context"
	"errors"
	"time"

	"github.com/enterprise-ai-platform/server/pkg/auth"
)

// Issuer 显示在认证器 App 中的发行方名称。
const Issuer = "政智通 GovAI"

// Store 2FA 持久化接口。
type Store interface {
	// GetStatus 返回是否启用与剩余可用备份码数量。
	GetStatus(ctx context.Context, userID string) (enabled bool, remaining int, err error)
	// GetSecret 返回用户的 TOTP 密钥（可能为空）。
	GetSecret(ctx context.Context, userID string) (secret string, err error)
	// GetSecretAndEnabled 返回密钥与启用状态。
	GetSecretAndEnabled(ctx context.Context, userID string) (secret string, enabled bool, err error)
	// SaveEnrollment 原子写入新密钥并重置备份码（未启用）。
	SaveEnrollment(ctx context.Context, userID, secret string, codeHashes []string) error
	// Enable 置 totp_enabled=true。
	Enable(ctx context.Context, userID string) error
	// Disable 关闭 2FA：清除密钥并删除所有备份码。
	Disable(ctx context.Context, userID string) error
	// ConsumeBackupCode 校验并一次性消费备份码，命中返回 true。
	ConsumeBackupCode(ctx context.Context, userID, code string) (bool, error)
}

var (
	ErrAlreadyEnabled = errors.New("两步验证已启用")
	ErrNotSetup       = errors.New("尚未开始两步验证设置")
	ErrBadCode        = errors.New("验证码或备份码错误")
)

// EnrollResult 是开始设置 2FA 的返回。
type EnrollResult struct {
	Secret      string
	OtpauthURI  string
	BackupCodes []string // 明文，仅返回一次
}

type Service struct {
	store       Store
	backupCount int
	nowUnix     func() int64
}

func NewService(store Store) *Service {
	return &Service{store: store, backupCount: 8, nowUnix: func() int64 { return time.Now().Unix() }}
}

// Status 返回当前 2FA 状态。
func (s *Service) Status(ctx context.Context, userID string) (enabled bool, remaining int, err error) {
	return s.store.GetStatus(ctx, userID)
}

// Enroll 生成新密钥与备份码并落库（未启用，需 Verify 确认）。
func (s *Service) Enroll(ctx context.Context, userID, email string) (*EnrollResult, error) {
	_, enabled, err := s.store.GetSecretAndEnabled(ctx, userID)
	if err != nil {
		return nil, err
	}
	if enabled {
		return nil, ErrAlreadyEnabled
	}

	secret, err := auth.GenerateTOTPSecret()
	if err != nil {
		return nil, err
	}
	plain, hashes, err := auth.GenerateBackupCodes(s.backupCount)
	if err != nil {
		return nil, err
	}
	if err := s.store.SaveEnrollment(ctx, userID, secret, hashes); err != nil {
		return nil, err
	}
	return &EnrollResult{
		Secret:      secret,
		OtpauthURI:  auth.TOTPProvisioningURI(secret, email, Issuer),
		BackupCodes: plain,
	}, nil
}

// EnableAfterVerify 校验首个验证码并启用 2FA。
func (s *Service) EnableAfterVerify(ctx context.Context, userID, code string) error {
	secret, err := s.store.GetSecret(ctx, userID)
	if err != nil {
		return err
	}
	if secret == "" {
		return ErrNotSetup
	}
	if !auth.ValidateTOTP(secret, code, s.nowUnix()) {
		return ErrBadCode
	}
	return s.store.Enable(ctx, userID)
}

// Disable 校验 TOTP 或备份码后关闭 2FA。未启用时视为成功（幂等）。
func (s *Service) Disable(ctx context.Context, userID, code, backupCode string) error {
	secret, enabled, err := s.store.GetSecretAndEnabled(ctx, userID)
	if err != nil {
		return err
	}
	if !enabled {
		return nil
	}
	if !s.verify(ctx, userID, secret, code, backupCode) {
		return ErrBadCode
	}
	return s.store.Disable(ctx, userID)
}

// VerifyLogin 在登录流程中校验 2FA：先试 TOTP，再试备份码（一次性消费）。
// secret 由调用方在登录查询时一并取出，避免重复查库。
func (s *Service) VerifyLogin(ctx context.Context, userID, secret, totpCode, backupCode string) bool {
	return s.verify(ctx, userID, secret, totpCode, backupCode)
}

func (s *Service) verify(ctx context.Context, userID, secret, totpCode, backupCode string) bool {
	if totpCode != "" && secret != "" && auth.ValidateTOTP(secret, totpCode, s.nowUnix()) {
		return true
	}
	if backupCode != "" {
		ok, err := s.store.ConsumeBackupCode(ctx, userID, backupCode)
		if err == nil && ok {
			return true
		}
	}
	return false
}