path: root/internal/controlplane/access.go

package controlplane

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"net/mail"
	"strings"
	"time"

	"aigw/internal/security"

	"github.com/jackc/pgx/v5"
)

var (
	ErrConsoleUnauthorized = errors.New("invalid console credentials")
	ErrConsoleRateLimited  = errors.New("too many login attempts")
)

const (
	RolePlatformAdmin   = "platform_admin"
	RolePlatformViewer  = "platform_viewer"
	RoleTenantAdmin     = "tenant_admin"
	RoleTenantBilling   = "tenant_billing"
	RoleTenantDeveloper = "tenant_developer"
	RoleTenantViewer    = "tenant_viewer"
)

func (a ConsoleActor) IsPlatform() bool { return strings.HasPrefix(a.Role, "platform_") }

func (a ConsoleActor) Can(permission string) bool {
	if a.Role == RolePlatformAdmin {
		return true
	}
	read := strings.HasSuffix(permission, ".read")
	switch a.Role {
	case RolePlatformViewer:
		return read
	case RoleTenantAdmin:
		switch permission {
		case "overview.read", "tenants.read", "projects.read", "projects.write", "keys.read", "keys.write",
			"billing.read", "billing.topup", "usage.read", "audit.read", "limits.read", "limits.write", "users.read", "users.write",
			"preferences.read", "developer.preferences.write", "billing.preferences.write":
			return true
		}
		return false
	case RoleTenantBilling:
		return permission == "overview.read" || permission == "billing.read" || permission == "billing.topup" || permission == "usage.read" || permission == "audit.read" || permission == "preferences.read" || permission == "billing.preferences.write"
	case RoleTenantDeveloper:
		return permission == "overview.read" || permission == "tenants.read" || permission == "projects.read" || permission == "keys.read" || permission == "keys.write" || permission == "usage.read" || permission == "limits.read" || permission == "preferences.read" || permission == "developer.preferences.write"
	case RoleTenantViewer:
		return permission == "overview.read" || permission == "tenants.read" || permission == "projects.read" || permission == "keys.read" || permission == "billing.read" || permission == "usage.read" || permission == "limits.read" || permission == "audit.read" || permission == "preferences.read"
	default:
		return false
	}
}

func (a ConsoleActor) Permissions() []string {
	all := []string{"overview.read", "preferences.read", "developer.preferences.write", "billing.preferences.write", "tenants.read", "tenants.write", "projects.read", "projects.write", "keys.read", "keys.write", "platform.read", "platform.write", "billing.read", "billing.topup", "billing.adjust", "usage.read", "limits.read", "limits.write", "users.read", "users.write", "audit.read"}
	result := make([]string, 0, len(all))
	for _, permission := range all {
		if a.Can(permission) {
			result = append(result, permission)
		}
	}
	return result
}

func (s *Store) AuthenticateConsoleToken(ctx context.Context, raw string) (ConsoleActor, error) {
	hash := sha256.Sum256([]byte(raw))
	var actor ConsoleActor
	err := s.db.QueryRow(ctx, `
		UPDATE console_users SET last_used_at = now()
		WHERE token_hash = $1 AND status = 'active'
		RETURNING id::text, COALESCE(tenant_id::text, ''), email, display_name, role`, hash[:],
	).Scan(&actor.ID, &actor.TenantID, &actor.Email, &actor.DisplayName, &actor.Role)
	if errors.Is(err, pgx.ErrNoRows) {
		return ConsoleActor{}, ErrConsoleUnauthorized
	}
	if err != nil {
		return ConsoleActor{}, fmt.Errorf("authenticate console token: %w", err)
	}
	return actor, nil
}

func (s *Store) ListConsoleUsers(ctx context.Context, tenantID string) ([]ConsoleUser, error) {
	query := `SELECT id::text, COALESCE(tenant_id::text, ''), email, display_name, role,
		COALESCE(token_prefix, ''), password_hash IS NOT NULL, status, email_verified_at IS NOT NULL,
		ARRAY_REMOVE(ARRAY[
			CASE WHEN EXISTS(SELECT 1 FROM console_totp_credentials tc WHERE tc.user_id=console_users.id AND tc.confirmed_at IS NOT NULL) THEN 'totp' END,
			CASE WHEN EXISTS(SELECT 1 FROM console_passkeys pk WHERE pk.user_id=console_users.id) THEN 'passkey' END
		],NULL), last_used_at, created_at FROM console_users`
	args := []any{}
	if tenantID != "" {
		query += ` WHERE tenant_id = $1`
		args = append(args, tenantID)
	}
	query += ` ORDER BY created_at DESC`
	rows, err := s.db.Query(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("query console users: %w", err)
	}
	defer rows.Close()
	result := make([]ConsoleUser, 0)
	for rows.Next() {
		var item ConsoleUser
		if err := rows.Scan(&item.ID, &item.TenantID, &item.Email, &item.DisplayName, &item.Role, &item.TokenPrefix, &item.HasPassword,
			&item.Status, &item.EmailVerified, &item.MFAMethods, &item.LastUsedAt, &item.CreatedAt); err != nil {
			return nil, fmt.Errorf("scan console user: %w", err)
		}
		result = append(result, item)
	}
	return result, rows.Err()
}

func (s *Store) CreateConsoleUser(ctx context.Context, input CreateConsoleUserInput) (CreatedConsoleUser, error) {
	input.Email = strings.ToLower(strings.TrimSpace(input.Email))
	input.DisplayName = strings.TrimSpace(input.DisplayName)
	input.TenantID = strings.TrimSpace(input.TenantID)
	address, addressErr := mail.ParseAddress(input.Email)
	if addressErr != nil || address.Address != input.Email || input.DisplayName == "" {
		return CreatedConsoleUser{}, errors.New("console user requires a valid email and display_name")
	}
	platform := input.Role == RolePlatformAdmin || input.Role == RolePlatformViewer
	tenant := input.Role == RoleTenantAdmin || input.Role == RoleTenantBilling || input.Role == RoleTenantDeveloper || input.Role == RoleTenantViewer
	if (!platform && !tenant) || (platform && input.TenantID != "") || (tenant && input.TenantID == "") {
		return CreatedConsoleUser{}, errors.New("console user role and tenant_id are inconsistent")
	}
	hash, salt, iterations, err := security.HashPassword(input.Password)
	if err != nil {
		return CreatedConsoleUser{}, err
	}
	var result CreatedConsoleUser
	err = s.db.QueryRow(ctx, `
		INSERT INTO console_users (tenant_id, email, display_name, role, password_hash, password_salt, password_iterations, password_changed_at)
		VALUES (NULLIF($1,'')::uuid,$2,$3,$4,$5,$6,$7,now())
		RETURNING id::text, COALESCE(tenant_id::text, ''), email, display_name, role, COALESCE(token_prefix,''),
			password_hash IS NOT NULL, status, last_used_at, created_at`,
		input.TenantID, input.Email, input.DisplayName, input.Role, hash, salt, iterations,
	).Scan(&result.ID, &result.TenantID, &result.Email, &result.DisplayName, &result.Role, &result.TokenPrefix,
		&result.HasPassword, &result.Status, &result.LastUsedAt, &result.CreatedAt)
	if err != nil {
		return CreatedConsoleUser{}, fmt.Errorf("create console user: %w", err)
	}
	return result, nil
}

func (s *Store) RevokeConsoleUser(ctx context.Context, id, tenantID string) error {
	tx, err := s.db.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)
	query := `UPDATE console_users SET status='revoked', revoked_at=now() WHERE id=$1 AND status='active'`
	args := []any{id}
	if tenantID != "" {
		query += ` AND tenant_id=$2`
		args = append(args, tenantID)
	}
	result, err := tx.Exec(ctx, query, args...)
	if err != nil {
		return err
	}
	if result.RowsAffected() == 0 {
		return ErrNotFound
	}
	if _, err := tx.Exec(ctx, `UPDATE console_sessions SET revoked_at=now() WHERE user_id=$1 AND revoked_at IS NULL`, id); err != nil {
		return fmt.Errorf("revoke console user sessions: %w", err)
	}
	return tx.Commit(ctx)
}

func (s *Store) RegisterTenant(ctx context.Context, input RegisterInput, currency string) (ConsoleActor, int64, error) {
	input.Organization = strings.TrimSpace(input.Organization)
	input.TenantSlug = strings.ToLower(strings.TrimSpace(input.TenantSlug))
	input.DisplayName = strings.TrimSpace(input.DisplayName)
	input.Email = strings.ToLower(strings.TrimSpace(input.Email))
	if input.Organization == "" || input.DisplayName == "" || !slugPattern.MatchString(input.TenantSlug) || !validEmail(input.Email) {
		return ConsoleActor{}, 0, errors.New("registration requires organization, a valid tenant_slug, display_name, and email")
	}
	if len(currency) != 3 {
		return ConsoleActor{}, 0, errors.New("registration currency is invalid")
	}
	hash, salt, iterations, err := security.HashPassword(input.Password)
	if err != nil {
		return ConsoleActor{}, 0, err
	}
	tx, err := s.db.Begin(ctx)
	if err != nil {
		return ConsoleActor{}, 0, err
	}
	defer tx.Rollback(ctx)
	var tenantID string
	if err := tx.QueryRow(ctx, `INSERT INTO tenants (slug,name) VALUES ($1,$2) RETURNING id::text`, input.TenantSlug, input.Organization).Scan(&tenantID); err != nil {
		return ConsoleActor{}, 0, fmt.Errorf("create registered tenant: %w", err)
	}
	if _, err := tx.Exec(ctx, `INSERT INTO projects (tenant_id,slug,name) VALUES ($1,'default','Default project')`, tenantID); err != nil {
		return ConsoleActor{}, 0, fmt.Errorf("create default project: %w", err)
	}
	if _, err := tx.Exec(ctx, `INSERT INTO tenant_wallets (tenant_id,currency) VALUES ($1,$2)`, tenantID, strings.ToLower(currency)); err != nil {
		return ConsoleActor{}, 0, fmt.Errorf("create tenant wallet: %w", err)
	}
	actor := ConsoleActor{TenantID: tenantID, Email: input.Email, DisplayName: input.DisplayName, Role: RoleTenantAdmin}
	if err := tx.QueryRow(ctx, `INSERT INTO console_users
		(tenant_id,email,display_name,role,password_hash,password_salt,password_iterations,password_changed_at)
		VALUES ($1,$2,$3,$4,$5,$6,$7,now()) RETURNING id::text`, tenantID, input.Email, input.DisplayName,
		RoleTenantAdmin, hash, salt, iterations).Scan(&actor.ID); err != nil {
		return ConsoleActor{}, 0, fmt.Errorf("create tenant administrator: %w", err)
	}
	generation, err := bumpGeneration(ctx, tx)
	if err != nil {
		return ConsoleActor{}, 0, err
	}
	if err := tx.Commit(ctx); err != nil {
		return ConsoleActor{}, 0, err
	}
	return actor, generation, nil
}

func (s *Store) AuthenticateConsolePassword(ctx context.Context, input PasswordLoginInput, remoteIP string) (ConsoleActor, error) {
	email := strings.ToLower(strings.TrimSpace(input.Email))
	emailIdentity := sha256.Sum256([]byte("email\x00" + email))
	ipIdentity := sha256.Sum256([]byte("ip\x00" + remoteIP))
	for _, identity := range [][]byte{emailIdentity[:], ipIdentity[:]} {
		var lockedUntil *time.Time
		err := s.db.QueryRow(ctx, `SELECT locked_until FROM console_login_throttles WHERE identity_hash=$1`, identity).Scan(&lockedUntil)
		if err != nil && !errors.Is(err, pgx.ErrNoRows) {
			return ConsoleActor{}, fmt.Errorf("read console login throttle: %w", err)
		}
		if lockedUntil != nil && lockedUntil.After(time.Now()) {
			return ConsoleActor{}, ErrConsoleRateLimited
		}
	}

	var actor ConsoleActor
	var expectedHash, salt []byte
	var iterations int
	err := s.db.QueryRow(ctx, `SELECT u.id::text, COALESCE(u.tenant_id::text,''), u.email, u.display_name, u.role,
		u.password_hash, u.password_salt, u.password_iterations
		FROM console_users u LEFT JOIN tenants t ON t.id=u.tenant_id
		WHERE lower(u.email)=$1 AND u.status='active' AND u.password_hash IS NOT NULL
		AND (u.tenant_id IS NULL OR t.status='active')`, email,
	).Scan(&actor.ID, &actor.TenantID, &actor.Email, &actor.DisplayName, &actor.Role, &expectedHash, &salt, &iterations)
	if errors.Is(err, pgx.ErrNoRows) {
		dummyHash := make([]byte, security.PasswordHashBytes)
		dummySalt := make([]byte, security.PasswordSaltBytes)
		_ = security.VerifyPassword(input.Password, dummyHash, dummySalt, security.PasswordIterations)
		if failureErr := s.recordLoginFailures(ctx, emailIdentity[:], ipIdentity[:]); failureErr != nil {
			return ConsoleActor{}, failureErr
		}
		return ConsoleActor{}, ErrConsoleUnauthorized
	}
	if err != nil {
		return ConsoleActor{}, fmt.Errorf("query console login: %w", err)
	}
	if !security.VerifyPassword(input.Password, expectedHash, salt, iterations) {
		if failureErr := s.recordLoginFailures(ctx, emailIdentity[:], ipIdentity[:]); failureErr != nil {
			return ConsoleActor{}, failureErr
		}
		return ConsoleActor{}, ErrConsoleUnauthorized
	}
	if _, err := s.db.Exec(ctx, `DELETE FROM console_login_throttles WHERE identity_hash=$1`, emailIdentity[:]); err != nil {
		return ConsoleActor{}, fmt.Errorf("clear console login throttle: %w", err)
	}
	if _, err := s.db.Exec(ctx, `UPDATE console_users SET last_used_at=now() WHERE id=$1`, actor.ID); err != nil {
		return ConsoleActor{}, fmt.Errorf("update console login time: %w", err)
	}
	actor.MFAMethods, err = s.MFAMethods(ctx, actor.ID)
	if err != nil {
		return ConsoleActor{}, err
	}
	return actor, nil
}

func (s *Store) recordLoginFailures(ctx context.Context, emailIdentity, ipIdentity []byte) error {
	if err := s.recordLoginFailure(ctx, emailIdentity, 5); err != nil {
		return err
	}
	return s.recordLoginFailure(ctx, ipIdentity, 25)
}

func (s *Store) recordLoginFailure(ctx context.Context, identityHash []byte, threshold int) error {
	_, err := s.db.Exec(ctx, `INSERT INTO console_login_throttles (identity_hash,failures) VALUES ($1,1)
		ON CONFLICT (identity_hash) DO UPDATE SET
		failures=CASE WHEN console_login_throttles.window_started_at < now()-interval '15 minutes' THEN 1 ELSE console_login_throttles.failures+1 END,
		window_started_at=CASE WHEN console_login_throttles.window_started_at < now()-interval '15 minutes' THEN now() ELSE console_login_throttles.window_started_at END,
		locked_until=CASE WHEN (CASE WHEN console_login_throttles.window_started_at < now()-interval '15 minutes' THEN 1 ELSE console_login_throttles.failures+1 END) >= $2 THEN now()+interval '15 minutes' ELSE NULL END,
		updated_at=now()`, identityHash, threshold)
	if err != nil {
		return fmt.Errorf("record console login failure: %w", err)
	}
	return nil
}

func (s *Store) CreateConsoleSession(ctx context.Context, actor ConsoleActor, ttl time.Duration, remoteIP, userAgent string) (ConsoleSession, error) {
	return s.CreateConsoleSessionWithMethod(ctx, actor, ttl, remoteIP, userAgent, "password", false)
}

func (s *Store) CreateConsoleSessionWithMethod(ctx context.Context, actor ConsoleActor, ttl time.Duration, remoteIP, userAgent, authMethod string, mfaVerified bool) (ConsoleSession, error) {
	if actor.ID == "" || ttl < time.Minute {
		return ConsoleSession{}, errors.New("session user and ttl are required")
	}
	token, tokenHash, err := randomCredential("sess-aigw-")
	if err != nil {
		return ConsoleSession{}, err
	}
	csrf, csrfHash, err := randomCredential("")
	if err != nil {
		return ConsoleSession{}, err
	}
	expiresAt := time.Now().UTC().Add(ttl)
	var sessionID string
	var mfaVerifiedAt any
	if mfaVerified {
		mfaVerifiedAt = time.Now().UTC()
	}
	err = s.db.QueryRow(ctx, `INSERT INTO console_sessions
		(user_id,token_hash,csrf_hash,expires_at,remote_ip,user_agent,auth_method,mfa_verified_at)
		VALUES ($1,$2,$3,$4,NULLIF($5,'')::inet,$6,$7,$8) RETURNING id::text`,
		actor.ID, tokenHash, csrfHash, expiresAt, remoteIP, userAgent, authMethod, mfaVerifiedAt).Scan(&sessionID)
	if err != nil {
		return ConsoleSession{}, fmt.Errorf("create console session: %w", err)
	}
	return ConsoleSession{Actor: actor, Token: token, CSRFToken: csrf, ExpiresAt: expiresAt, ID: sessionID, AuthMethod: authMethod}, nil
}

func (s *Store) AuthenticateConsoleSession(ctx context.Context, rawToken string) (ConsoleActor, []byte, error) {
	hash := sha256.Sum256([]byte(rawToken))
	var actor ConsoleActor
	var csrfHash []byte
	err := s.db.QueryRow(ctx, `UPDATE console_sessions s SET last_seen_at=CASE WHEN s.last_seen_at < now()-interval '5 minutes' THEN now() ELSE s.last_seen_at END
		FROM console_users u LEFT JOIN tenants t ON t.id=u.tenant_id
		WHERE s.user_id=u.id AND s.token_hash=$1 AND s.revoked_at IS NULL AND s.expires_at>now()
		AND u.status='active' AND (u.tenant_id IS NULL OR t.status='active')
		RETURNING u.id::text,COALESCE(u.tenant_id::text,''),u.email,u.display_name,u.role,s.csrf_hash`, hash[:],
	).Scan(&actor.ID, &actor.TenantID, &actor.Email, &actor.DisplayName, &actor.Role, &csrfHash)
	if errors.Is(err, pgx.ErrNoRows) {
		return ConsoleActor{}, nil, ErrConsoleUnauthorized
	}
	if err != nil {
		return ConsoleActor{}, nil, fmt.Errorf("authenticate console session: %w", err)
	}
	actor.MFAMethods, err = s.MFAMethods(ctx, actor.ID)
	if err != nil {
		return ConsoleActor{}, nil, err
	}
	return actor, csrfHash, nil
}

func (s *Store) RevokeConsoleSession(ctx context.Context, rawToken string) error {
	if rawToken == "" {
		return nil
	}
	hash := sha256.Sum256([]byte(rawToken))
	_, err := s.db.Exec(ctx, `UPDATE console_sessions SET revoked_at=now() WHERE token_hash=$1 AND revoked_at IS NULL`, hash[:])
	return err
}

func (s *Store) ListDeviceSessions(ctx context.Context, actorID, currentToken string) ([]DeviceSession, error) {
	currentHash := sha256.Sum256([]byte(currentToken))
	rows, err := s.db.Query(ctx, `SELECT id::text,COALESCE(host(remote_ip),''),user_agent,auth_method,mfa_verified_at,
		created_at,last_seen_at,expires_at,token_hash=$2 FROM console_sessions
		WHERE user_id=$1 AND revoked_at IS NULL AND expires_at>now() ORDER BY last_seen_at DESC`, actorID, currentHash[:])
	if err != nil {
		return nil, fmt.Errorf("list console sessions: %w", err)
	}
	defer rows.Close()
	result := make([]DeviceSession, 0)
	for rows.Next() {
		var item DeviceSession
		if err := rows.Scan(&item.ID, &item.RemoteIP, &item.UserAgent, &item.AuthMethod, &item.MFAVerifiedAt,
			&item.CreatedAt, &item.LastSeenAt, &item.ExpiresAt, &item.Current); err != nil {
			return nil, err
		}
		result = append(result, item)
	}
	return result, rows.Err()
}

func (s *Store) RevokeDeviceSession(ctx context.Context, actorID, sessionID string) (bool, error) {
	tag, err := s.db.Exec(ctx, `UPDATE console_sessions SET revoked_at=now()
		WHERE id=$1 AND user_id=$2 AND revoked_at IS NULL`, sessionID, actorID)
	if err != nil {
		return false, err
	}
	if tag.RowsAffected() == 0 {
		return false, ErrNotFound
	}
	return true, nil
}

func (s *Store) RevokeOtherDeviceSessions(ctx context.Context, actorID, currentToken string) error {
	currentHash := sha256.Sum256([]byte(currentToken))
	_, err := s.db.Exec(ctx, `UPDATE console_sessions SET revoked_at=now()
		WHERE user_id=$1 AND token_hash<>$2 AND revoked_at IS NULL`, actorID, currentHash[:])
	return err
}

func (s *Store) ChangeConsolePassword(ctx context.Context, actorID string, input PasswordChangeInput) error {
	var expectedHash, salt []byte
	var iterations int
	if err := s.db.QueryRow(ctx, `SELECT password_hash,password_salt,password_iterations FROM console_users WHERE id=$1 AND status='active'`, actorID).
		Scan(&expectedHash, &salt, &iterations); err != nil {
		return ErrConsoleUnauthorized
	}
	if !security.VerifyPassword(input.CurrentPassword, expectedHash, salt, iterations) {
		return ErrConsoleUnauthorized
	}
	hash, newSalt, newIterations, err := security.HashPassword(input.NewPassword)
	if err != nil {
		return err
	}
	tx, err := s.db.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)
	if _, err := tx.Exec(ctx, `UPDATE console_users SET password_hash=$2,password_salt=$3,password_iterations=$4,password_changed_at=now() WHERE id=$1`, actorID, hash, newSalt, newIterations); err != nil {
		return err
	}
	if _, err := tx.Exec(ctx, `UPDATE console_sessions SET revoked_at=now() WHERE user_id=$1 AND revoked_at IS NULL`, actorID); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func randomCredential(prefix string) (string, []byte, error) {
	random := make([]byte, 32)
	if _, err := rand.Read(random); err != nil {
		return "", nil, err
	}
	raw := prefix + base64.RawURLEncoding.EncodeToString(random)
	hash := sha256.Sum256([]byte(raw))
	return raw, hash[:], nil
}

func validEmail(value string) bool {
	address, err := mail.ParseAddress(value)
	return err == nil && address.Address == value
}