notifier/internal/auth/keystore_db.go

package auth

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"time"

	"github.com/igodwin/notifier/internal/logging"
	"github.com/lib/pq"
)

// KeyStoreDB provides persistent storage for API keys using PostgreSQL.
// It acts as the backend for the in-memory cache. Only SHA-256 digests of
// keys are persisted — the raw secret never reaches the database.
type KeyStoreDB struct {
	db     *sql.DB
	logger *logging.Logger
}

// NewKeyStoreDB creates a new database-backed key store
func NewKeyStoreDB(dbURL string, logger *logging.Logger) (*KeyStoreDB, error) {
	db, err := sql.Open("postgres", dbURL)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to database: %w", err)
	}

	db.SetMaxOpenConns(10)
	db.SetMaxIdleConns(5)
	db.SetConnMaxLifetime(30 * time.Minute)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	if err := db.PingContext(ctx); err != nil {
		return nil, fmt.Errorf("failed to ping database: %w", err)
	}

	ks := &KeyStoreDB{db: db, logger: logger}

	if err := ks.initializeSchema(); err != nil {
		return nil, fmt.Errorf("failed to initialize schema: %w", err)
	}

	return ks, nil
}

// initializeSchema creates the necessary tables and indexes if they don't
// exist, and migrates legacy plaintext-key rows to hashed storage.
func (ks *KeyStoreDB) initializeSchema() error {
	tableSchema := `
	-- API Keys table (key_hash is the SHA-256 digest of the raw key)
	CREATE TABLE IF NOT EXISTS api_keys (
		id SERIAL PRIMARY KEY,
		key_hash VARCHAR(64) UNIQUE NOT NULL,
		key_preview VARCHAR(16) NOT NULL DEFAULT '',
		name VARCHAR(255) NOT NULL,
		client_id VARCHAR(255) NOT NULL,
		roles TEXT[] DEFAULT '{}',
		created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
		last_used_at TIMESTAMP,
		expires_at TIMESTAMP,
		is_active BOOLEAN NOT NULL DEFAULT true,
		rate_limit INTEGER NOT NULL DEFAULT 0,
		created_by VARCHAR(255),
		metadata JSONB DEFAULT '{}'::jsonb
	);

	-- Audit log for key operations
	CREATE TABLE IF NOT EXISTS api_key_audit_log (
		id SERIAL PRIMARY KEY,
		key_id INTEGER NOT NULL REFERENCES api_keys(id),
		action VARCHAR(50) NOT NULL,
		performed_by VARCHAR(255) NOT NULL,
		performed_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
		details JSONB DEFAULT '{}'::jsonb
	);
	`

	if _, err := ks.db.Exec(tableSchema); err != nil {
		return fmt.Errorf("failed to create tables: %w", err)
	}

	if err := ks.migrateLegacyPlaintextKeys(); err != nil {
		return fmt.Errorf("failed to migrate legacy plaintext keys: %w", err)
	}

	indexSchema := `
	CREATE INDEX IF NOT EXISTS idx_api_keys_client_id ON api_keys(client_id);
	CREATE INDEX IF NOT EXISTS idx_api_keys_active ON api_keys(is_active);
	CREATE INDEX IF NOT EXISTS idx_api_keys_expires ON api_keys(expires_at);
	CREATE INDEX IF NOT EXISTS idx_api_key_audit_log_key_id ON api_key_audit_log(key_id);
	CREATE INDEX IF NOT EXISTS idx_api_key_audit_log_performed_at ON api_key_audit_log(performed_at);
	`

	if _, err := ks.db.Exec(indexSchema); err != nil {
		return fmt.Errorf("failed to create indexes: %w", err)
	}

	return nil
}

// migrateLegacyPlaintextKeys upgrades tables created by earlier versions that
// stored the raw key in a "key" column: it adds the hash columns, hashes each
// plaintext key in place, then drops the plaintext column entirely.
func (ks *KeyStoreDB) migrateLegacyPlaintextKeys() error {
	var hasLegacyColumn bool
	err := ks.db.QueryRow(`
		SELECT EXISTS (
			SELECT 1 FROM information_schema.columns
			WHERE table_name = 'api_keys' AND column_name = 'key'
		)`).Scan(&hasLegacyColumn)
	if err != nil {
		return fmt.Errorf("failed to inspect schema: %w", err)
	}
	if !hasLegacyColumn {
		return nil
	}

	if _, err := ks.db.Exec(`
		ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS key_hash VARCHAR(64);
		ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS key_preview VARCHAR(16) NOT NULL DEFAULT '';
	`); err != nil {
		return fmt.Errorf("failed to add hash columns: %w", err)
	}

	rows, err := ks.db.Query(`SELECT id, key FROM api_keys WHERE key_hash IS NULL AND key IS NOT NULL`)
	if err != nil {
		return fmt.Errorf("failed to read legacy keys: %w", err)
	}
	defer rows.Close()

	type legacyRow struct {
		id  int
		key string
	}
	var legacy []legacyRow
	for rows.Next() {
		var r legacyRow
		if err := rows.Scan(&r.id, &r.key); err != nil {
			return fmt.Errorf("failed to scan legacy key: %w", err)
		}
		legacy = append(legacy, r)
	}
	if err := rows.Err(); err != nil {
		return err
	}

	for _, r := range legacy {
		if _, err := ks.db.Exec(
			`UPDATE api_keys SET key_hash = $1, key_preview = $2 WHERE id = $3`,
			HashKey(r.key), keyPreview(r.key), r.id,
		); err != nil {
			return fmt.Errorf("failed to hash legacy key id=%d: %w", r.id, err)
		}
	}

	// Drop the plaintext column and enforce uniqueness on the digest.
	if _, err := ks.db.Exec(`
		ALTER TABLE api_keys DROP COLUMN key;
		ALTER TABLE api_keys ALTER COLUMN key_hash SET NOT NULL;
		CREATE UNIQUE INDEX IF NOT EXISTS idx_api_keys_key_hash ON api_keys(key_hash);
	`); err != nil {
		return fmt.Errorf("failed to finalize hash migration: %w", err)
	}

	if ks.logger != nil {
		ks.logger.Infof("Migrated %d legacy plaintext API key(s) to hashed storage", len(legacy))
	}
	return nil
}

// SaveKey persists an API key to the database. Only the digest and preview
// are stored; apiKey.Key (the raw secret) is never written.
func (ks *KeyStoreDB) SaveKey(ctx context.Context, key *APIKey, createdBy string) error {
	if key.KeyHash == "" {
		return fmt.Errorf("refusing to save key without digest")
	}

	query := `
	INSERT INTO api_keys (key_hash, key_preview, name, client_id, roles, created_at, last_used_at, expires_at, is_active, rate_limit, created_by)
	VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
	ON CONFLICT (key_hash) DO UPDATE SET
		name = EXCLUDED.name,
		roles = EXCLUDED.roles,
		is_active = EXCLUDED.is_active,
		rate_limit = EXCLUDED.rate_limit,
		last_used_at = COALESCE(EXCLUDED.last_used_at, api_keys.last_used_at)
	`

	_, err := ks.db.ExecContext(ctx, query,
		key.KeyHash,
		key.KeyPreview,
		key.Name,
		key.ClientID,
		pq.Array(key.Roles),
		key.CreatedAt,
		key.LastUsedAt,
		key.ExpiresAt,
		key.IsActive,
		key.RateLimit,
		createdBy,
	)

	if err != nil {
		return fmt.Errorf("failed to save key: %w", err)
	}

	ks.logAudit(ctx, key.KeyHash, "created", createdBy, map[string]interface{}{
		"client_id": key.ClientID,
		"roles":     key.Roles,
	})

	return nil
}

// scanKey scans a single api_keys row into an APIKey.
func scanKey(scanner interface{ Scan(...interface{}) error }) (*APIKey, error) {
	var key APIKey
	var roles []string

	err := scanner.Scan(
		&key.KeyHash,
		&key.KeyPreview,
		&key.Name,
		&key.ClientID,
		pq.Array(&roles),
		&key.CreatedAt,
		&key.LastUsedAt,
		&key.ExpiresAt,
		&key.IsActive,
		&key.RateLimit,
	)
	if err != nil {
		return nil, err
	}

	key.Roles = roles
	return &key, nil
}

const keyColumns = `key_hash, key_preview, name, client_id, roles, created_at, last_used_at, expires_at, is_active, rate_limit`

// GetKeyByHash retrieves an API key from the database by its digest
func (ks *KeyStoreDB) GetKeyByHash(ctx context.Context, keyHash string) (*APIKey, error) {
	row := ks.db.QueryRowContext(ctx,
		`SELECT `+keyColumns+` FROM api_keys WHERE key_hash = $1`, keyHash)

	key, err := scanKey(row)
	if err == sql.ErrNoRows {
		return nil, ErrKeyNotFound
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get key: %w", err)
	}
	return key, nil
}

// ListKeys retrieves all active API keys for a client
func (ks *KeyStoreDB) ListKeys(ctx context.Context, clientID string) ([]*APIKey, error) {
	rows, err := ks.db.QueryContext(ctx,
		`SELECT `+keyColumns+` FROM api_keys
		 WHERE client_id = $1 AND is_active = true
		 ORDER BY created_at DESC`, clientID)
	if err != nil {
		return nil, fmt.Errorf("failed to list keys: %w", err)
	}
	defer rows.Close()

	var keys []*APIKey
	for rows.Next() {
		key, err := scanKey(rows)
		if err != nil {
			return nil, fmt.Errorf("failed to scan key: %w", err)
		}
		keys = append(keys, key)
	}

	return keys, rows.Err()
}

// DeactivateKeyByHash disables an API key identified by its digest
func (ks *KeyStoreDB) DeactivateKeyByHash(ctx context.Context, keyHash string, deactivatedBy string) error {
	result, err := ks.db.ExecContext(ctx,
		`UPDATE api_keys SET is_active = false WHERE key_hash = $1`, keyHash)
	if err != nil {
		return fmt.Errorf("failed to deactivate key: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return err
	}

	if rows == 0 {
		return ErrKeyNotFound
	}

	ks.logAudit(ctx, keyHash, "deactivated", deactivatedBy, nil)
	return nil
}

// UpdateLastUsed updates the last_used_at timestamp for a key digest
func (ks *KeyStoreDB) UpdateLastUsed(ctx context.Context, keyHash string) error {
	_, err := ks.db.ExecContext(ctx,
		`UPDATE api_keys SET last_used_at = CURRENT_TIMESTAMP WHERE key_hash = $1`, keyHash)
	if err != nil {
		return fmt.Errorf("failed to update last used: %w", err)
	}
	return nil
}

// LoadAllKeys loads all active keys into memory for caching
func (ks *KeyStoreDB) LoadAllKeys(ctx context.Context) ([]*APIKey, error) {
	rows, err := ks.db.QueryContext(ctx,
		`SELECT `+keyColumns+` FROM api_keys
		 WHERE is_active = true AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)`)
	if err != nil {
		return nil, fmt.Errorf("failed to load keys: %w", err)
	}
	defer rows.Close()

	var keys []*APIKey
	for rows.Next() {
		key, err := scanKey(rows)
		if err != nil {
			return nil, fmt.Errorf("failed to scan key: %w", err)
		}
		keys = append(keys, key)
	}

	return keys, rows.Err()
}

// logAudit logs a key operation to the audit trail. Failures don't abort the
// calling operation but are logged rather than silently dropped.
func (ks *KeyStoreDB) logAudit(ctx context.Context, keyHash string, action string, performedBy string, details map[string]interface{}) {
	var keyID int
	if err := ks.db.QueryRowContext(ctx,
		"SELECT id FROM api_keys WHERE key_hash = $1", keyHash).Scan(&keyID); err != nil {
		ks.warnf("audit log: failed to resolve key for action=%s: %v", action, err)
		return
	}

	detailsJSON := []byte("{}")
	if len(details) > 0 {
		encoded, err := json.Marshal(details)
		if err != nil {
			ks.warnf("audit log: failed to encode details for action=%s: %v", action, err)
		} else {
			detailsJSON = encoded
		}
	}

	if _, err := ks.db.ExecContext(ctx,
		`INSERT INTO api_key_audit_log (key_id, action, performed_by, details) VALUES ($1, $2, $3, $4)`,
		keyID, action, performedBy, detailsJSON); err != nil {
		ks.warnf("audit log: failed to record action=%s: %v", action, err)
	}
}

func (ks *KeyStoreDB) warnf(format string, args ...interface{}) {
	if ks.logger != nil {
		ks.logger.Warnf(format, args...)
	}
}

// Ping verifies database connectivity (used by readiness checks).
func (ks *KeyStoreDB) Ping(ctx context.Context) error {
	return ks.db.PingContext(ctx)
}

// Close closes the database connection
func (ks *KeyStoreDB) Close() error {
	return ks.db.Close()
}

// GetAuditLog retrieves audit log entries for a key digest
func (ks *KeyStoreDB) GetAuditLog(ctx context.Context, keyHash string, limit int) ([]map[string]interface{}, error) {
	return ks.auditLogQuery(ctx,
		`SELECT al.action, al.performed_by, al.performed_at, al.details
		 FROM api_key_audit_log al
		 JOIN api_keys ak ON al.key_id = ak.id
		 WHERE ak.key_hash = $1
		 ORDER BY al.performed_at DESC
		 LIMIT $2`, keyHash, limit)
}

// GetKeyByName retrieves an API key by its name
func (ks *KeyStoreDB) GetKeyByName(ctx context.Context, name string) (*APIKey, error) {
	row := ks.db.QueryRowContext(ctx,
		`SELECT `+keyColumns+` FROM api_keys WHERE name = $1`, name)

	key, err := scanKey(row)
	if err == sql.ErrNoRows {
		return nil, ErrKeyNotFound
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get key by name: %w", err)
	}
	return key, nil
}

// DeactivateKeyByName disables an API key by its name
func (ks *KeyStoreDB) DeactivateKeyByName(ctx context.Context, name string, deactivatedBy string) error {
	key, err := ks.GetKeyByName(ctx, name)
	if err != nil {
		return err
	}

	return ks.DeactivateKeyByHash(ctx, key.KeyHash, deactivatedBy)
}

// GetAuditLogByName retrieves audit log entries for a key identified by name
func (ks *KeyStoreDB) GetAuditLogByName(ctx context.Context, name string, limit int) ([]map[string]interface{}, error) {
	return ks.auditLogQuery(ctx,
		`SELECT al.action, al.performed_by, al.performed_at, al.details
		 FROM api_key_audit_log al
		 JOIN api_keys ak ON al.key_id = ak.id
		 WHERE ak.name = $1
		 ORDER BY al.performed_at DESC
		 LIMIT $2`, name, limit)
}

func (ks *KeyStoreDB) auditLogQuery(ctx context.Context, query string, ident string, limit int) ([]map[string]interface{}, error) {
	rows, err := ks.db.QueryContext(ctx, query, ident, limit)
	if err != nil {
		return nil, fmt.Errorf("failed to get audit log: %w", err)
	}
	defer rows.Close()

	var logs []map[string]interface{}
	for rows.Next() {
		var action, performedBy, details string
		var performedAt time.Time

		if err := rows.Scan(&action, &performedBy, &performedAt, &details); err != nil {
			return nil, err
		}

		logs = append(logs, map[string]interface{}{
			"action":       action,
			"performed_by": performedBy,
			"performed_at": performedAt,
			"details":      details,
		})
	}

	return logs, rows.Err()
}