archivmail/internal/userstore/userstore.go

package userstore

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"golang.org/x/crypto/bcrypt"
)

const (
	RoleUser         = "user"
	RoleAdmin        = "admin" // legacy, maps to domain_admin
	RoleAuditor      = "auditor"
	RoleDomainAdmin  = "domain_admin"
	RoleDomainAuditor = "domain_auditor"
	RoleSuperAdmin   = "superadmin"

	bcryptCost = 12
)

// User represents a user account in the system.
type User struct {
	ID          int64      `json:"id"`
	Username    string     `json:"username"`
	Email       string     `json:"email"`
	Role        string     `json:"role"`
	Source      string     `json:"source"` // "local" or "ldap"
	Active      bool       `json:"active"`
	CreatedAt   time.Time  `json:"created_at"`
	TenantID    *int64     `json:"tenant_id,omitempty"`
	TOTPEnabled bool       `json:"totp_enabled"`
	TOTPResetAt *time.Time `json:"totp_reset_at,omitempty"`
	TOTPResetBy *string    `json:"totp_reset_by,omitempty"`
}

// CreateUserRequest holds parameters for creating a new user.
type CreateUserRequest struct {
	Username string
	Email    string
	Password string
	Role     string
	TenantID *int64
}

// UpdateUserRequest holds optional fields for updating a user.
type UpdateUserRequest struct {
	Email    *string
	Role     *string
	Active   *bool
	Password *string
}

// Store is a PostgreSQL-backed user store.
type Store struct {
	pool *pgxpool.Pool
}

// New connects to PostgreSQL using the given DSN and initialises the schema.
func New(dsn string) (*Store, error) {
	ctx := context.Background()
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		return nil, fmt.Errorf("userstore: connect: %w", err)
	}

	s := &Store{pool: pool}
	if err := s.initSchema(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("userstore: init schema: %w", err)
	}

	return s, nil
}

func (s *Store) initSchema(ctx context.Context) error {
	_, err := s.pool.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS users (
			id              BIGSERIAL PRIMARY KEY,
			username        VARCHAR(100) UNIQUE NOT NULL,
			email           VARCHAR(255) UNIQUE NOT NULL,
			password_hash   VARCHAR(255) NOT NULL DEFAULT '',
			role            VARCHAR(20)  NOT NULL DEFAULT 'user',
			source          VARCHAR(20)  NOT NULL DEFAULT 'local',
			active          BOOLEAN      NOT NULL DEFAULT true,
			created_at      TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
			last_login_at   TIMESTAMPTZ
		);
		ALTER TABLE users ADD COLUMN IF NOT EXISTS last_login_at TIMESTAMPTZ;
		ALTER TABLE users ADD COLUMN IF NOT EXISTS tenant_id BIGINT;
		CREATE TABLE IF NOT EXISTS token_blacklist (
			jti        VARCHAR(255) PRIMARY KEY,
			expires_at TIMESTAMPTZ  NOT NULL
		);
		CREATE TABLE IF NOT EXISTS login_attempts (
			username    VARCHAR(100) NOT NULL,
			ip          VARCHAR(45)  NOT NULL,
			attempted_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
		);
		CREATE INDEX IF NOT EXISTS idx_login_attempts_username_time ON login_attempts (username, attempted_at);
		CREATE INDEX IF NOT EXISTS idx_users_tenant ON users (tenant_id);
	`)
	if err != nil {
		return err
	}
	// PROJ-24: TOTP 2FA columns
	_, err = s.pool.Exec(ctx, `
		ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_secret   BYTEA;
		ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_enabled  BOOLEAN NOT NULL DEFAULT false;
		ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_reset_at TIMESTAMPTZ;
		ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_reset_by TEXT;
	`)
	return err
}

// Close closes the underlying connection pool.
func (s *Store) Close() error {
	s.pool.Close()
	return nil
}

// Create inserts a new local user with a bcrypt-hashed password.
func (s *Store) Create(req CreateUserRequest) (*User, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(req.Password), bcryptCost)
	if err != nil {
		return nil, fmt.Errorf("userstore: bcrypt: %w", err)
	}

	ctx := context.Background()
	var id int64
	err = s.pool.QueryRow(ctx,
		`INSERT INTO users (username, email, password_hash, role, source, active, created_at, tenant_id)
		 VALUES ($1, $2, $3, $4, 'local', true, NOW(), $5)
		 RETURNING id`,
		req.Username, req.Email, string(hash), req.Role, req.TenantID,
	).Scan(&id)
	if err != nil {
		return nil, fmt.Errorf("userstore: create: %w", err)
	}

	return s.GetByID(id)
}

// GetByID retrieves a user by their numeric ID.
func (s *Store) GetByID(id int64) (*User, error) {
	ctx := context.Background()
	row := s.pool.QueryRow(ctx,
		`SELECT id, username, email, role, source, active, created_at, tenant_id, totp_enabled, totp_reset_at, totp_reset_by FROM users WHERE id = $1`, id,
	)
	return scanUser(row)
}

// GetByUsername retrieves a user by their username.
func (s *Store) GetByUsername(username string) (*User, error) {
	ctx := context.Background()
	row := s.pool.QueryRow(ctx,
		`SELECT id, username, email, role, source, active, created_at, tenant_id, totp_enabled, totp_reset_at, totp_reset_by FROM users WHERE username = $1`, username,
	)
	return scanUser(row)
}

// VerifyPassword checks credentials and returns the user, or an error if the
// password is wrong or the account is disabled.
func (s *Store) VerifyPassword(username, password string) (*User, error) {
	ctx := context.Background()
	row := s.pool.QueryRow(ctx,
		`SELECT id, username, email, role, source, active, created_at, tenant_id, totp_enabled, totp_reset_at, totp_reset_by, password_hash FROM users WHERE username = $1`,
		username,
	)

	var u User
	var hash string
	err := row.Scan(&u.ID, &u.Username, &u.Email, &u.Role, &u.Source, &u.Active, &u.CreatedAt, &u.TenantID, &u.TOTPEnabled, &u.TOTPResetAt, &u.TOTPResetBy, &hash)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, errors.New("userstore: user not found")
	}
	if err != nil {
		return nil, fmt.Errorf("userstore: scan: %w", err)
	}

	if !u.Active {
		return nil, errors.New("userstore: account disabled")
	}

	if err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(password)); err != nil {
		return nil, errors.New("userstore: wrong password")
	}

	return &u, nil
}

// Update applies a partial update to a user record.
func (s *Store) Update(id int64, req UpdateUserRequest) (*User, error) {
	ctx := context.Background()

	if req.Email != nil {
		if _, err := s.pool.Exec(ctx, `UPDATE users SET email = $1 WHERE id = $2`, *req.Email, id); err != nil {
			return nil, fmt.Errorf("userstore: update email: %w", err)
		}
	}
	if req.Role != nil {
		if _, err := s.pool.Exec(ctx, `UPDATE users SET role = $1 WHERE id = $2`, *req.Role, id); err != nil {
			return nil, fmt.Errorf("userstore: update role: %w", err)
		}
	}
	if req.Active != nil {
		if _, err := s.pool.Exec(ctx, `UPDATE users SET active = $1 WHERE id = $2`, *req.Active, id); err != nil {
			return nil, fmt.Errorf("userstore: update active: %w", err)
		}
	}
	if req.Password != nil {
		hash, err := bcrypt.GenerateFromPassword([]byte(*req.Password), bcryptCost)
		if err != nil {
			return nil, fmt.Errorf("userstore: bcrypt: %w", err)
		}
		if _, err := s.pool.Exec(ctx, `UPDATE users SET password_hash = $1 WHERE id = $2`, string(hash), id); err != nil {
			return nil, fmt.Errorf("userstore: update password: %w", err)
		}
	}

	return s.GetByID(id)
}

// Delete removes a user by ID. Returns an error if the user does not exist.
func (s *Store) Delete(id int64) error {
	ctx := context.Background()
	tag, err := s.pool.Exec(ctx, `DELETE FROM users WHERE id = $1`, id)
	if err != nil {
		return fmt.Errorf("userstore: delete: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("userstore: user %d not found", id)
	}
	return nil
}

// List returns all users, optionally filtered by role. Pass role="" to list all.
func (s *Store) List(role string) ([]*User, error) {
	ctx := context.Background()
	var rows pgx.Rows
	var err error

	if role == "" {
		rows, err = s.pool.Query(ctx,
			`SELECT id, username, email, role, source, active, created_at, tenant_id, totp_enabled, totp_reset_at, totp_reset_by FROM users ORDER BY id`)
	} else {
		rows, err = s.pool.Query(ctx,
			`SELECT id, username, email, role, source, active, created_at, tenant_id, totp_enabled, totp_reset_at, totp_reset_by FROM users WHERE role = $1 ORDER BY id`, role)
	}
	if err != nil {
		return nil, fmt.Errorf("userstore: list: %w", err)
	}
	defer rows.Close()

	var users []*User
	for rows.Next() {
		u, err := scanUserRow(rows)
		if err != nil {
			return nil, err
		}
		users = append(users, u)
	}
	return users, rows.Err()
}

// ListByTenant returns all users belonging to a specific tenant.
func (s *Store) ListByTenant(ctx context.Context, tenantID int64) ([]*User, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT id, username, email, role, source, active, created_at, tenant_id, totp_enabled, totp_reset_at, totp_reset_by FROM users WHERE tenant_id = $1 ORDER BY id`,
		tenantID,
	)
	if err != nil {
		return nil, fmt.Errorf("userstore: list by tenant: %w", err)
	}
	defer rows.Close()

	var users []*User
	for rows.Next() {
		u, err := scanUserRow(rows)
		if err != nil {
			return nil, err
		}
		users = append(users, u)
	}
	return users, rows.Err()
}

// BlacklistToken adds a JWT ID to the token blacklist.
func (s *Store) BlacklistToken(jti string, expires time.Time) error {
	ctx := context.Background()
	_, err := s.pool.Exec(ctx,
		`INSERT INTO token_blacklist (jti, expires_at) VALUES ($1, $2)
		 ON CONFLICT (jti) DO UPDATE SET expires_at = EXCLUDED.expires_at`,
		jti, expires.UTC(),
	)
	return err
}

// IsBlacklisted returns true if the given JTI is in the blacklist.
func (s *Store) IsBlacklisted(jti string) (bool, error) {
	ctx := context.Background()
	var count int
	err := s.pool.QueryRow(ctx,
		`SELECT COUNT(*) FROM token_blacklist WHERE jti = $1`, jti,
	).Scan(&count)
	return count > 0, err
}

// UpdateLastLogin sets last_login_at to now for the given user.
func (s *Store) UpdateLastLogin(id int64) error {
	ctx := context.Background()
	_, err := s.pool.Exec(ctx, `UPDATE users SET last_login_at = NOW() WHERE id = $1`, id)
	return err
}

// RecordLoginAttempt inserts a failed login attempt record.
func (s *Store) RecordLoginAttempt(username, ip string) error {
	ctx := context.Background()
	_, err := s.pool.Exec(ctx,
		`INSERT INTO login_attempts (username, ip, attempted_at) VALUES ($1, $2, NOW())`,
		username, ip,
	)
	return err
}

// CountRecentFailures returns the number of failed attempts for username in the last window.
func (s *Store) CountRecentFailures(username string, window time.Duration) (int, error) {
	ctx := context.Background()
	var count int
	err := s.pool.QueryRow(ctx,
		`SELECT COUNT(*) FROM login_attempts WHERE username = $1 AND attempted_at > NOW() - $2::interval`,
		username, window.String(),
	).Scan(&count)
	return count, err
}

// AdminCount returns the number of active privileged users (admin, domain_admin, superadmin).
func (s *Store) AdminCount() (int, error) {
	ctx := context.Background()
	var count int
	err := s.pool.QueryRow(ctx,
		`SELECT COUNT(*) FROM users WHERE role IN ('admin','domain_admin','superadmin') AND active = true`,
	).Scan(&count)
	return count, err
}

// DeleteSafe removes a user but refuses if they are the last active admin.
func (s *Store) DeleteSafe(id int64) error {
	user, err := s.GetByID(id)
	if err != nil {
		return err
	}
	if user.Role == RoleAdmin || user.Role == RoleDomainAdmin || user.Role == RoleSuperAdmin {
		count, err := s.AdminCount()
		if err != nil {
			return fmt.Errorf("userstore: admin count: %w", err)
		}
		if count <= 1 {
			return fmt.Errorf("userstore: cannot delete last admin")
		}
	}
	return s.Delete(id)
}

// CleanExpiredTokens removes blacklist entries whose expiry has passed.
func (s *Store) CleanExpiredTokens() error {
	ctx := context.Background()
	_, err := s.pool.Exec(ctx, `DELETE FROM token_blacklist WHERE expires_at < NOW()`)
	return err
}

// UpsertLDAPUser creates or updates an LDAP-sourced user.
// tenantID may be nil for users not associated with a specific tenant.
func (s *Store) UpsertLDAPUser(username, email, role string, tenantID *int64) (*User, error) {
	ctx := context.Background()
	_, err := s.pool.Exec(ctx, `
		INSERT INTO users (username, email, password_hash, role, source, active, created_at, tenant_id)
		VALUES ($1, $2, '', $3, 'ldap', true, NOW(), $4)
		ON CONFLICT (username) DO UPDATE SET
			email     = EXCLUDED.email,
			role      = EXCLUDED.role,
			source    = 'ldap',
			tenant_id = COALESCE(EXCLUDED.tenant_id, users.tenant_id)
	`, username, email, role, tenantID)
	if err != nil {
		return nil, fmt.Errorf("userstore: upsert ldap: %w", err)
	}
	return s.GetByUsername(username)
}

// --- helpers ---

func scanUser(row pgx.Row) (*User, error) {
	var u User
	err := row.Scan(&u.ID, &u.Username, &u.Email, &u.Role, &u.Source, &u.Active, &u.CreatedAt, &u.TenantID, &u.TOTPEnabled, &u.TOTPResetAt, &u.TOTPResetBy)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, fmt.Errorf("userstore: not found")
	}
	if err != nil {
		return nil, fmt.Errorf("userstore: scan: %w", err)
	}
	return &u, nil
}

func scanUserRow(rows pgx.Rows) (*User, error) {
	var u User
	if err := rows.Scan(&u.ID, &u.Username, &u.Email, &u.Role, &u.Source, &u.Active, &u.CreatedAt, &u.TenantID, &u.TOTPEnabled, &u.TOTPResetAt, &u.TOTPResetBy); err != nil {
		return nil, fmt.Errorf("userstore: scan row: %w", err)
	}
	return &u, nil
}

// ── PROJ-24: TOTP 2FA Methods ────────────────────────────────────────────

// SetTOTPSecret stores the encrypted TOTP secret (not yet activated).
func (s *Store) SetTOTPSecret(ctx context.Context, userID int64, encryptedSecret []byte) error {
	_, err := s.pool.Exec(ctx, `UPDATE users SET totp_secret = $1 WHERE id = $2`, encryptedSecret, userID)
	if err != nil {
		return fmt.Errorf("userstore: set totp secret: %w", err)
	}
	return nil
}

// EnableTOTP activates TOTP for the user (after code confirmation).
func (s *Store) EnableTOTP(ctx context.Context, userID int64) error {
	_, err := s.pool.Exec(ctx, `UPDATE users SET totp_enabled = true WHERE id = $1`, userID)
	if err != nil {
		return fmt.Errorf("userstore: enable totp: %w", err)
	}
	return nil
}

// DisableTOTP deactivates TOTP and removes the secret (user self-service).
func (s *Store) DisableTOTP(ctx context.Context, userID int64) error {
	_, err := s.pool.Exec(ctx, `UPDATE users SET totp_enabled = false, totp_secret = NULL WHERE id = $1`, userID)
	if err != nil {
		return fmt.Errorf("userstore: disable totp: %w", err)
	}
	return nil
}

// ResetTOTP resets TOTP for a user (admin action) and logs who performed the reset.
func (s *Store) ResetTOTP(ctx context.Context, userID int64, resetBy string) error {
	_, err := s.pool.Exec(ctx,
		`UPDATE users SET totp_enabled = false, totp_secret = NULL, totp_reset_at = NOW(), totp_reset_by = $1 WHERE id = $2`,
		resetBy, userID,
	)
	if err != nil {
		return fmt.Errorf("userstore: reset totp: %w", err)
	}
	return nil
}

// ── PROJ-25: Profile Update Methods ───────────────────────────────────────

// GetPasswordHash returns the bcrypt password hash for a user by ID.
func (s *Store) GetPasswordHash(ctx context.Context, userID int64) (string, error) {
	var hash string
	err := s.pool.QueryRow(ctx, `SELECT password_hash FROM users WHERE id = $1`, userID).Scan(&hash)
	if err != nil {
		return "", fmt.Errorf("userstore: get password hash: %w", err)
	}
	return hash, nil
}

// UpdatePassword sets a new password hash for the given user.
func (s *Store) UpdatePassword(ctx context.Context, userID int64, passwordHash string) error {
	_, err := s.pool.Exec(ctx, `UPDATE users SET password_hash = $1 WHERE id = $2`, passwordHash, userID)
	if err != nil {
		return fmt.Errorf("userstore: update password: %w", err)
	}
	return nil
}

// UpdateEmail sets a new email address for the given user.
func (s *Store) UpdateEmail(ctx context.Context, userID int64, email string) error {
	_, err := s.pool.Exec(ctx, `UPDATE users SET email = $1 WHERE id = $2`, email, userID)
	if err != nil {
		return fmt.Errorf("userstore: update email: %w", err)
	}
	return nil
}

// GetTOTPSecret returns the encrypted TOTP secret and enabled status for a user.
func (s *Store) GetTOTPSecret(ctx context.Context, userID int64) (secret []byte, enabled bool, err error) {
	err = s.pool.QueryRow(ctx,
		`SELECT totp_secret, totp_enabled FROM users WHERE id = $1`, userID,
	).Scan(&secret, &enabled)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, false, fmt.Errorf("userstore: user not found")
	}
	if err != nil {
		return nil, false, fmt.Errorf("userstore: get totp secret: %w", err)
	}
	return secret, enabled, nil
}