ai-gateway-go/internal/identity/service.go

package identity

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	"aigateway.local/core/internal/platform/config"
	"aigateway.local/core/internal/platform/cryptox"
)

var (
	ErrInvalidCredentials = errors.New("invalid credentials")
	ErrAccountDisabled    = errors.New("account disabled")
	ErrInvalidTOTP        = errors.New("invalid or reused TOTP code")
	ErrTOTPAlreadyEnabled = errors.New("TOTP is already enabled")
	ErrTOTPNotEnabled     = errors.New("TOTP is not enabled")
	ErrTOTPSetupRequired  = errors.New("TOTP setup is required")
)

const dummyPasswordHash = "pbkdf2_sha256$600000$00112233445566778899aabbccddeeff$afca0887b188255f525e15e30f5aa5a0b210a3e253bfaf9630411f0782bb6573"

type LockedError struct {
	Until time.Time
}

func (e LockedError) Error() string {
	return fmt.Sprintf("account locked until %s", e.Until.Format(time.RFC3339))
}

type LoginResult struct {
	Token       string
	TempToken   string
	RequireTOTP bool
	Account     Account
}

type TOTPSetupResult struct {
	Secret          string
	ProvisioningURI string
}

type Service struct {
	repository                   *Repository
	sessions                     *SessionStore
	limiter                      *LoginLimiter
	hasher                       PasswordHasher
	config                       config.Auth
	totpCipher                   cryptox.Cipher
	idpCipher                    cryptox.Cipher
	allowPrivateIdentityProvider bool
	oidcClient                   *http.Client
	samlMetadataMu               sync.RWMutex
	samlMetadata                 map[string]samlMetadataCacheEntry
	now                          func() time.Time
}

func (s *Service) SetIdentityProviderCipher(cipher cryptox.Cipher, allowPrivate bool) {
	s.idpCipher = cipher
	s.allowPrivateIdentityProvider = allowPrivate
	s.oidcClient = newOIDCHTTPClient(allowPrivate)
}

func NewService(repository *Repository, sessions *SessionStore, limiter *LoginLimiter, cfg config.Auth, totpCipher cryptox.Cipher) *Service {
	return &Service{repository: repository, sessions: sessions, limiter: limiter, hasher: PasswordHasher{}, config: cfg, totpCipher: totpCipher, oidcClient: newOIDCHTTPClient(false), samlMetadata: make(map[string]samlMetadataCacheEntry), now: time.Now}
}

// AllowLogin reports whether a login attempt from ip may proceed. When the
// per-IP sliding-window limit is exceeded it returns false (caller responds 429).
func (s *Service) AllowLogin(ctx context.Context, ip string) bool {
	return s.limiter == nil || s.limiter.Allow(ctx, ip)
}

func (s *Service) Login(ctx context.Context, kind Kind, login, password string) (LoginResult, error) {
	account, err := s.findByLogin(ctx, kind, login)
	if errors.Is(err, ErrNotFound) {
		_ = s.hasher.Verify(password, dummyPasswordHash)
		return LoginResult{}, ErrInvalidCredentials
	}
	if err != nil {
		return LoginResult{}, err
	}
	if !account.Active {
		return LoginResult{}, ErrAccountDisabled
	}
	if account.Locked(s.now()) {
		return LoginResult{}, LockedError{Until: *account.LockedUntil}
	}
	if account.PasswordHash == "" || !s.hasher.Verify(password, account.PasswordHash) {
		lockedUntil, recordErr := s.repository.RecordFailure(ctx, account, s.config.MaxFailures, s.config.LockDuration)
		if recordErr != nil {
			return LoginResult{}, recordErr
		}
		if lockedUntil != nil && lockedUntil.After(s.now()) {
			return LoginResult{}, LockedError{Until: *lockedUntil}
		}
		return LoginResult{}, ErrInvalidCredentials
	}
	if account.TOTPEnabled {
		principal := principalFor(account)
		token, tokenErr := s.sessions.CreatePending(ctx, principal, s.config.TOTPChallengeTTL)
		if tokenErr != nil {
			return LoginResult{}, tokenErr
		}
		return LoginResult{TempToken: token, RequireTOTP: true, Account: account}, nil
	}

	var upgradedHash *string
	if s.hasher.NeedsUpgrade(account.PasswordHash) {
		hash, hashErr := s.hasher.Hash(password)
		if hashErr != nil {
			return LoginResult{}, hashErr
		}
		upgradedHash = &hash
	}
	principal := principalFor(account)
	token, err := s.sessions.Create(ctx, principal)
	if err != nil {
		return LoginResult{}, err
	}
	if err := s.repository.CompleteLogin(ctx, account, upgradedHash); err != nil {
		_ = s.sessions.Delete(ctx, "Bearer "+token)
		return LoginResult{}, err
	}
	return LoginResult{Token: token, Account: account}, nil
}

func (s *Service) CompleteTOTPLogin(ctx context.Context, kind Kind, tempToken, code, backupCode string) (LoginResult, error) {
	principal, err := s.sessions.AuthenticatePending(ctx, tempToken, kind)
	if err != nil {
		return LoginResult{}, err
	}
	account, err := s.findByID(ctx, kind, principal.SubjectID)
	if err != nil {
		return LoginResult{}, err
	}
	if !account.Active {
		return LoginResult{}, ErrAccountDisabled
	}
	if account.Locked(s.now()) {
		return LoginResult{}, LockedError{Until: *account.LockedUntil}
	}
	if !account.TOTPEnabled {
		return LoginResult{}, ErrTOTPNotEnabled
	}
	valid, err := s.verifyAndConsumeFactor(ctx, account, code, backupCode)
	if err != nil {
		return LoginResult{}, err
	}
	if !valid {
		lockedUntil, recordErr := s.repository.RecordFailure(ctx, account, s.config.MaxFailures, s.config.LockDuration)
		if recordErr != nil {
			return LoginResult{}, recordErr
		}
		if lockedUntil != nil && lockedUntil.After(s.now()) {
			return LoginResult{}, LockedError{Until: *lockedUntil}
		}
		return LoginResult{}, ErrInvalidTOTP
	}
	token, err := s.sessions.Create(ctx, principalFor(account))
	if err != nil {
		return LoginResult{}, err
	}
	if err := s.repository.CompleteLogin(ctx, account, nil); err != nil {
		_ = s.sessions.Delete(ctx, "Bearer "+token)
		return LoginResult{}, err
	}
	_ = s.sessions.DeleteToken(ctx, tempToken)
	return LoginResult{Token: token, Account: account}, nil
}

func (s *Service) SetupTOTP(ctx context.Context, account Account, password string) (TOTPSetupResult, error) {
	if account.TOTPEnabled {
		return TOTPSetupResult{}, ErrTOTPAlreadyEnabled
	}
	if !s.hasher.Verify(password, account.PasswordHash) {
		return TOTPSetupResult{}, ErrInvalidCredentials
	}
	secret, err := GenerateTOTPSecret()
	if err != nil {
		return TOTPSetupResult{}, err
	}
	encrypted, version, err := s.totpCipher.Encrypt([]byte(secret))
	if err != nil {
		return TOTPSetupResult{}, err
	}
	if err := s.repository.SetTOTPSecret(ctx, account, encrypted, version); err != nil {
		return TOTPSetupResult{}, err
	}
	return TOTPSetupResult{Secret: secret, ProvisioningURI: TOTPProvisioningURI(secret, account.Kind, account.Login)}, nil
}

func (s *Service) ConfirmTOTP(ctx context.Context, account Account, code string) ([]string, error) {
	account, err := s.findByID(ctx, account.Kind, account.ID)
	if err != nil {
		return nil, err
	}
	if account.TOTPEnabled {
		return nil, ErrTOTPAlreadyEnabled
	}
	secret, err := s.decryptTOTPSecret(account)
	if err != nil {
		return nil, err
	}
	step, valid := VerifyTOTP(secret, code, s.now())
	if !valid {
		return nil, ErrInvalidTOTP
	}
	codes, records, err := GenerateBackupCodes()
	if err != nil {
		return nil, err
	}
	if err := s.repository.EnableTOTP(ctx, account, step, records); err != nil {
		return nil, err
	}
	return codes, nil
}

func (s *Service) DisableTOTP(ctx context.Context, account Account, password, code, backupCode string) error {
	account, err := s.findByID(ctx, account.Kind, account.ID)
	if err != nil {
		return err
	}
	if !account.TOTPEnabled {
		return ErrTOTPNotEnabled
	}
	if !s.hasher.Verify(password, account.PasswordHash) {
		return ErrInvalidCredentials
	}
	valid, err := s.verifyAndConsumeFactor(ctx, account, code, backupCode)
	if err != nil {
		return err
	}
	if !valid {
		return ErrInvalidTOTP
	}
	return s.repository.DisableTOTP(ctx, account)
}

func (s *Service) RegenerateBackupCodes(ctx context.Context, account Account, password, code, backupCode string) ([]string, error) {
	account, err := s.findByID(ctx, account.Kind, account.ID)
	if err != nil {
		return nil, err
	}
	if !account.TOTPEnabled {
		return nil, ErrTOTPNotEnabled
	}
	if !s.hasher.Verify(password, account.PasswordHash) {
		return nil, ErrInvalidCredentials
	}
	valid, err := s.verifyAndConsumeFactor(ctx, account, code, backupCode)
	if err != nil {
		return nil, err
	}
	if !valid {
		return nil, ErrInvalidTOTP
	}
	codes, records, err := GenerateBackupCodes()
	if err != nil {
		return nil, err
	}
	if err := s.repository.ReplaceBackupCodes(ctx, account, records); err != nil {
		return nil, err
	}
	return codes, nil
}

func (s *Service) verifyAndConsumeFactor(ctx context.Context, account Account, code, backupCode string) (bool, error) {
	if strings.TrimSpace(backupCode) != "" {
		return s.repository.ConsumeBackupCode(ctx, account, HashBackupCode(backupCode))
	}
	secret, err := s.decryptTOTPSecret(account)
	if err != nil {
		return false, err
	}
	step, valid := VerifyTOTP(secret, code, s.now())
	if !valid {
		return false, nil
	}
	return s.repository.ConsumeTOTPStep(ctx, account, step)
}

func (s *Service) decryptTOTPSecret(account Account) (string, error) {
	if len(account.EncryptedTOTPSecret) == 0 || account.TOTPKekVersion == nil {
		return "", ErrTOTPSetupRequired
	}
	plaintext, err := s.totpCipher.Decrypt(account.EncryptedTOTPSecret, *account.TOTPKekVersion)
	if err != nil {
		return "", err
	}
	return string(plaintext), nil
}

func (s *Service) Authenticate(ctx context.Context, kind Kind, authorization string) (Account, error) {
	principal, err := s.sessions.Authenticate(ctx, authorization, kind)
	if err != nil {
		return Account{}, err
	}
	var account Account
	if kind == KindAdmin {
		account, err = s.repository.FindAdminByID(ctx, principal.SubjectID)
	} else {
		account, err = s.repository.FindPortalByID(ctx, principal.SubjectID)
	}
	if errors.Is(err, ErrNotFound) {
		return Account{}, ErrInvalidSession
	}
	if err != nil {
		return Account{}, err
	}
	if !account.Active {
		return Account{}, ErrAccountDisabled
	}
	return account, nil
}

func (s *Service) Logout(ctx context.Context, authorization string) error {
	return s.sessions.Delete(ctx, authorization)
}

// ChangePassword updates the authenticated account password after verifying the
// current password. Externally provisioned portal accounts without a local
// password may set their first password without an old-password check.
func (s *Service) ChangePassword(ctx context.Context, account Account, oldPassword, newPassword string) error {
	account, err := s.findByID(ctx, account.Kind, account.ID)
	if err != nil {
		return err
	}
	if account.PasswordHash != "" && !s.hasher.Verify(oldPassword, account.PasswordHash) {
		return ErrInvalidCredentials
	}
	if len(newPassword) < 12 || len(newPassword) > 1024 {
		return errors.New("new password must contain 12 to 1024 characters")
	}
	hash, err := s.hasher.Hash(newPassword)
	if err != nil {
		return err
	}
	return s.repository.SetPassword(ctx, account, hash)
}

func (s *Service) findByLogin(ctx context.Context, kind Kind, login string) (Account, error) {
	if kind == KindAdmin {
		return s.repository.FindAdminByLogin(ctx, login)
	}
	return s.repository.FindPortalByLogin(ctx, login)
}

func (s *Service) findByID(ctx context.Context, kind Kind, id string) (Account, error) {
	if kind == KindAdmin {
		return s.repository.FindAdminByID(ctx, id)
	}
	return s.repository.FindPortalByID(ctx, id)
}

func principalFor(account Account) Principal {
	return Principal{Kind: account.Kind, SubjectID: account.ID, Login: account.Login, DisplayName: account.DisplayName, Role: account.Role}
}