gochat/internal/auth/mfa.go

package auth

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base32"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math"
	"strings"
	"time"

	"github.com/gochat/gochat/internal/model"
	pkgcrypto "github.com/gochat/gochat/pkg/crypto"
	"gorm.io/gorm"
)

const mfaBackupCodesAttribute = "mfa_backup_code_hashes"

// Reference: P2E §1.5 — MFA (TOTP) support
// Implements time-based one-time password (TOTP) per RFC 6238.
// Corresponds to Chatwoot enterprise TwoFactorAuthController pattern.

// MFAService manages multi-factor authentication using TOTP.
type MFAService struct {
	db *gorm.DB
}

// NewMFAService creates a MFA service backed by GORM.
func NewMFAService(db *gorm.DB) *MFAService {
	return &MFAService{db: db}
}

// TOTPConfig holds TOTP algorithm parameters.
// Standard parameters per RFC 6238: SHA-1, 30-second step, 6 digits.
type TOTPConfig struct {
	Period    uint64 // time step in seconds (default: 30)
	Digits    int    // number of digits (default: 6)
	Algorithm string // hash algorithm (default: SHA1)
	Issuer    string // issuer name for QR code (default: GoChat)
}

// DefaultTOTPConfig returns standard TOTP parameters.
func DefaultTOTPConfig() TOTPConfig {
	return TOTPConfig{
		Period:    30,
		Digits:    6,
		Algorithm: "SHA1",
		Issuer:    "GoChat",
	}
}

// GenerateTOTPSecret creates a random base32-encoded TOTP secret for a user.
// The secret is 160 bits (20 bytes) per RFC 4226 recommendation.
func (s *MFAService) GenerateTOTPSecret(userID uint) (string, string, error) {
	// Generate 20 random bytes for 160-bit secret
	secretBytes := make([]byte, 20)
	if _, err := rand.Read(secretBytes); err != nil {
		return "", "", fmt.Errorf("failed to generate random secret: %w", err)
	}

	// Encode as base32 (uppercase, no padding per RFC 4648)
	secret := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(secretBytes)

	// Look up user email for QR code URI
	var user model.User
	if err := s.db.First(&user, userID).Error; err != nil {
		return "", "", fmt.Errorf("user not found: %w", err)
	}

	// Generate otpauth URI for QR code scanning
	cfg := DefaultTOTPConfig()
	uri := fmt.Sprintf("otpauth://totp/%s:%s?secret=%s&issuer=%s&algorithm=%s&digits=%d&period=%d",
		cfg.Issuer,
		user.Email,
		secret,
		cfg.Issuer,
		cfg.Algorithm,
		cfg.Digits,
		cfg.Period,
	)

	return secret, uri, nil
}

// EnableTOTP stores the TOTP secret for a user after successful verification.
// This is a two-step process: user must verify a TOTP code before enabling.
func (s *MFAService) EnableTOTP(userID uint, secret string) error {
	var user model.User
	if err := s.db.First(&user, userID).Error; err != nil {
		return fmt.Errorf("user not found: %w", err)
	}

	user.TOTPSecret = secret
	user.TOTPEnabled = true

	if err := s.db.Save(&user).Error; err != nil {
		return fmt.Errorf("failed to enable totp: %w", err)
	}

	return nil
}

// BeginTOTPSetup creates and stores a pending TOTP secret for Chatwoot's
// profile MFA setup flow. The user is activated only after VerifyAndActivateTOTP.
func (s *MFAService) BeginTOTPSetup(userID uint) (string, string, error) {
	secret, uri, err := s.GenerateTOTPSecret(userID)
	if err != nil {
		return "", "", err
	}

	var user model.User
	if err := s.db.First(&user, userID).Error; err != nil {
		return "", "", fmt.Errorf("user not found: %w", err)
	}
	user.TOTPSecret = secret
	user.TOTPEnabled = false
	if err := s.db.Save(&user).Error; err != nil {
		return "", "", fmt.Errorf("failed to store pending totp secret: %w", err)
	}
	return secret, uri, nil
}

// VerifyAndActivateTOTP validates the pending profile MFA code, enables MFA,
// and returns the one-time backup codes expected by Chatwoot's verify response.
func (s *MFAService) VerifyAndActivateTOTP(userID uint, code string) ([]string, error) {
	var user model.User
	if err := s.db.First(&user, userID).Error; err != nil {
		return nil, fmt.Errorf("user not found: %w", err)
	}
	if user.TOTPSecret == "" {
		return nil, fmt.Errorf("mfa setup not started for user")
	}
	if !validateTOTP(user.TOTPSecret, code, DefaultTOTPConfig()) {
		return nil, fmt.Errorf("invalid totp code")
	}
	user.TOTPEnabled = true
	if err := s.db.Save(&user).Error; err != nil {
		return nil, fmt.Errorf("failed to enable totp: %w", err)
	}
	return s.GenerateBackupCodes(userID)
}

// VerifyTOTPCode validates a TOTP code against the user's stored secret.
// Uses a 1-period window (±30 seconds) to account for clock drift per RFC 6238.
func (s *MFAService) VerifyTOTPCode(userID uint, code string) (bool, error) {
	var user model.User
	if err := s.db.First(&user, userID).Error; err != nil {
		return false, fmt.Errorf("user not found: %w", err)
	}

	if !user.TOTPEnabled || user.TOTPSecret == "" {
		return false, fmt.Errorf("mfa not enabled for user")
	}

	cfg := DefaultTOTPConfig()
	return validateTOTP(user.TOTPSecret, code, cfg), nil
}

// DisableTOTP removes TOTP configuration for a user.
// Requires verification of current TOTP code before disabling.
func (s *MFAService) DisableTOTP(userID uint, code string) error {
	// Verify current code before allowing disable
	valid, err := s.VerifyTOTPCode(userID, code)
	if err != nil {
		return fmt.Errorf("verification failed: %w", err)
	}
	if !valid {
		return fmt.Errorf("invalid totp code")
	}

	var user model.User
	if err := s.db.First(&user, userID).Error; err != nil {
		return fmt.Errorf("user not found: %w", err)
	}

	user.TOTPSecret = ""
	user.TOTPEnabled = false

	if err := s.db.Save(&user).Error; err != nil {
		return fmt.Errorf("failed to disable totp: %w", err)
	}

	return nil
}

// DisableTOTPWithPassword mirrors Chatwoot profile MFA destroy: the current
// password and either an OTP code or a backup code must be provided.
func (s *MFAService) DisableTOTPWithPassword(userID uint, password, code, backupCode string) error {
	var user model.User
	if err := s.db.First(&user, userID).Error; err != nil {
		return fmt.Errorf("user not found: %w", err)
	}
	if !user.TOTPEnabled || user.TOTPSecret == "" {
		return fmt.Errorf("mfa not enabled for user")
	}
	if !pkgcrypto.CheckPassword(password, user.PasswordDigest) && !pkgcrypto.CheckPassword(password, user.Password) {
		return fmt.Errorf("invalid credentials")
	}
	if backupCode != "" {
		if err := s.consumeBackupCode(&user, backupCode); err != nil {
			return err
		}
	} else if !validateTOTP(user.TOTPSecret, code, DefaultTOTPConfig()) {
		return fmt.Errorf("invalid totp code")
	}

	user.TOTPSecret = ""
	user.TOTPEnabled = false
	if err := s.db.Save(&user).Error; err != nil {
		return fmt.Errorf("failed to disable totp: %w", err)
	}
	return nil
}

// IsMFAEnabled checks whether MFA is enabled for a user.
func (s *MFAService) IsMFAEnabled(userID uint) (bool, error) {
	var user model.User
	if err := s.db.Select("totp_enabled").First(&user, userID).Error; err != nil {
		return false, fmt.Errorf("user not found: %w", err)
	}
	return user.TOTPEnabled, nil
}

// BackupCodesGenerated reports whether the user currently has MFA backup codes.
func (s *MFAService) BackupCodesGenerated(userID uint) (bool, error) {
	var user model.User
	if err := s.db.Select("custom_attributes").First(&user, userID).Error; err != nil {
		return false, fmt.Errorf("user not found: %w", err)
	}
	codes := backupCodeHashes(user.CustomAttributes)
	return len(codes) > 0, nil
}

// validateTOTP validates a TOTP code against a secret using the given config.
// Allows ±1 period window for clock drift tolerance.
func validateTOTP(secret string, code string, cfg TOTPConfig) bool {
	now := time.Now().Unix()
	period := int64(cfg.Period)

	// Check current period and ±1 period for clock drift
	for offset := -1; offset <= 1; offset++ {
		t := (now + int64(offset)*period) / period
		expected := generateTOTP(secret, t, cfg)
		if expected == strings.TrimSpace(code) {
			return true
		}
	}

	return false
}

// generateTOTP generates a TOTP code for the given time counter.
// Implements RFC 6238 algorithm: HMAC-SHA1 with time-based counter.
func generateTOTP(secret string, timeCounter int64, cfg TOTPConfig) string {
	// Decode base32 secret
	key, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(strings.ToUpper(secret))
	if err != nil {
		return ""
	}

	// Encode time counter as 8-byte big-endian
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, uint64(timeCounter))

	// HMAC-SHA1
	h := hmac.New(sha1.New, key)
	h.Write(buf)
	hash := h.Sum(nil)

	// Dynamic truncation per RFC 4226
	offset := hash[len(hash)-1] & 0x0f
	truncated := (int32(hash[offset]&0x7f) << 24) |
		(int32(hash[offset+1]&0xff) << 16) |
		(int32(hash[offset+2]&0xff) << 8) |
		(int32(hash[offset+3] & 0xff))

	// Modulo 10^digits
	mod := int32(math.Pow10(cfg.Digits))
	code := truncated % mod

	// Format with leading zeros to achieve correct digit count
	return fmt.Sprintf("%0*d", cfg.Digits, code)
}

// ValidateTOTPCode is a public helper that validates a TOTP code against a given secret.
// Used by handlers to verify TOTP codes during MFA setup before enabling on the user.
func ValidateTOTPCode(secret string, code string, cfg TOTPConfig) bool {
	return validateTOTP(secret, code, cfg)
}

// --- math.Pow10 helper ---
func init() {
	// Ensure math package is linked
	_ = math.E
}

// GenerateBackupCodes creates a set of one-time backup codes for MFA recovery.
// Reference: Chatwoot MfaController#backup_codes
func (s *MFAService) GenerateBackupCodes(userID uint) ([]string, error) {
	var user model.User
	if err := s.db.First(&user, userID).Error; err != nil {
		return nil, fmt.Errorf("user not found: %w", err)
	}
	var codes []string
	var hashes []string
	for i := 0; i < 10; i++ {
		code := cryptoRandomString(8)
		codes = append(codes, code)
		hash, err := pkgcrypto.HashPassword(code)
		if err != nil {
			return nil, fmt.Errorf("failed to hash backup code: %w", err)
		}
		hashes = append(hashes, hash)
	}
	attrs := customAttributesMap(user.CustomAttributes)
	attrs[mfaBackupCodesAttribute] = hashes
	encoded, err := json.Marshal(attrs)
	if err != nil {
		return nil, fmt.Errorf("failed to encode backup codes: %w", err)
	}
	user.CustomAttributes = encoded
	if err := s.db.Save(&user).Error; err != nil {
		return nil, fmt.Errorf("failed to store backup codes: %w", err)
	}
	return codes, nil
}

func (s *MFAService) consumeBackupCode(user *model.User, code string) error {
	hashes := backupCodeHashes(user.CustomAttributes)
	for i, hash := range hashes {
		if pkgcrypto.CheckPassword(code, hash) {
			hashes = append(hashes[:i], hashes[i+1:]...)
			attrs := customAttributesMap(user.CustomAttributes)
			attrs[mfaBackupCodesAttribute] = hashes
			encoded, err := json.Marshal(attrs)
			if err != nil {
				return fmt.Errorf("failed to encode backup codes: %w", err)
			}
			user.CustomAttributes = encoded
			if err := s.db.Save(user).Error; err != nil {
				return fmt.Errorf("failed to consume backup code: %w", err)
			}
			return nil
		}
	}
	return fmt.Errorf("invalid backup code")
}

func backupCodeHashes(raw []byte) []string {
	attrs := customAttributesMap(raw)
	value, ok := attrs[mfaBackupCodesAttribute]
	if !ok {
		return nil
	}
	items, ok := value.([]any)
	if !ok {
		return nil
	}
	hashes := make([]string, 0, len(items))
	for _, item := range items {
		if text, ok := item.(string); ok && text != "" {
			hashes = append(hashes, text)
		}
	}
	return hashes
}

func customAttributesMap(raw []byte) map[string]any {
	attrs := map[string]any{}
	if len(raw) == 0 {
		return attrs
	}
	_ = json.Unmarshal(raw, &attrs)
	return attrs
}

// cryptoRandomString generates a random alphanumeric string of given length.
func cryptoRandomString(length int) string {
	const charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	b := make([]byte, length)
	for i := range b {
		buf := make([]byte, 1)
		_, _ = rand.Read(buf)
		b[i] = charset[int(buf[0])%len(charset)]
	}
	return string(b)
}