fetch_ml/internal/auth/database.go

package auth

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

	_ "github.com/mattn/go-sqlite3" // SQLite driver
)

// DatabaseAuthStore implements authentication using SQLite database
type DatabaseAuthStore struct {
	db *sql.DB
}

// APIKeyRecord represents an API key in the database
type APIKeyRecord struct {
	CreatedAt   time.Time  `json:"created_at"`
	ExpiresAt   *time.Time `json:"expires_at,omitempty"`
	RevokedAt   *time.Time `json:"revoked_at,omitempty"`
	UserID      string     `json:"user_id"`
	KeyHash     string     `json:"key_hash"`
	Roles       string     `json:"roles"`
	Permissions string     `json:"permissions"`
	ID          int        `json:"id"`
	Admin       bool       `json:"admin"`
}

// NewDatabaseAuthStore creates a new database-backed auth store
func NewDatabaseAuthStore(dbPath string) (*DatabaseAuthStore, error) {
	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}

	store := &DatabaseAuthStore{db: db}
	if err := store.init(); err != nil {
		return nil, fmt.Errorf("failed to initialize database: %w", err)
	}

	return store, nil
}

// init creates the necessary database tables
func (s *DatabaseAuthStore) init() error {
	ctx := context.Background()
	query := `
	CREATE TABLE IF NOT EXISTS api_keys (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		user_id TEXT NOT NULL UNIQUE,
		key_hash TEXT NOT NULL UNIQUE,
		admin BOOLEAN NOT NULL DEFAULT FALSE,
		roles TEXT NOT NULL DEFAULT '[]',
		permissions TEXT NOT NULL DEFAULT '{}',
		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		expires_at DATETIME,
		revoked_at DATETIME,
		CHECK (json_valid(roles)),
		CHECK (json_valid(permissions))
	);

	CREATE INDEX IF NOT EXISTS idx_api_keys_hash ON api_keys(key_hash);
	CREATE INDEX IF NOT EXISTS idx_api_keys_user ON api_keys(user_id);
	CREATE INDEX IF NOT EXISTS idx_api_keys_active ON api_keys(
		revoked_at,
		COALESCE(expires_at, '9999-12-31')
	);
	`

	_, err := s.db.ExecContext(ctx, query)
	return err
}

// ValidateAPIKey checks if an API key is valid and returns user info
func (s *DatabaseAuthStore) ValidateAPIKey(ctx context.Context, key string) (*User, error) {
	keyHash := HashAPIKey(key)

	query := `
	SELECT user_id, admin, roles, permissions, expires_at, revoked_at
	FROM api_keys
	WHERE key_hash = ? 
	  AND (revoked_at IS NULL OR revoked_at > ?)
	  AND (expires_at IS NULL OR expires_at > ?)
	`

	var userID string
	var admin bool
	var rolesJSON, permissionsJSON string
	var expiresAt, revokedAt sql.NullTime
	now := time.Now()

	err := s.db.QueryRowContext(ctx, query, keyHash, now, now).Scan(
		&userID, &admin, &rolesJSON, &permissionsJSON, &expiresAt, &revokedAt,
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, fmt.Errorf("invalid API key")
		}
		return nil, fmt.Errorf("database error: %w", err)
	}

	// Parse roles
	var roles []string
	if err := json.Unmarshal([]byte(rolesJSON), &roles); err != nil {
		log.Printf("Failed to parse roles for user %s: %v", userID, err)
		roles = []string{}
	}

	// Parse permissions
	var permissions map[string]bool
	if err := json.Unmarshal([]byte(permissionsJSON), &permissions); err != nil {
		log.Printf("Failed to parse permissions for user %s: %v", userID, err)
		permissions = make(map[string]bool)
	}

	// Admin gets all permissions
	if admin {
		permissions["*"] = true
	}

	return &User{
		Name:        userID,
		Admin:       admin,
		Roles:       roles,
		Permissions: permissions,
	}, nil
}

// CreateAPIKey creates a new API key in the database
func (s *DatabaseAuthStore) CreateAPIKey(
	ctx context.Context,
	userID string,
	keyHash string,
	admin bool,
	roles []string,
	permissions map[string]bool,
	expiresAt *time.Time,
) error {
	rolesJSON, err := json.Marshal(roles)
	if err != nil {
		return fmt.Errorf("failed to marshal roles: %w", err)
	}

	permissionsJSON, err := json.Marshal(permissions)
	if err != nil {
		return fmt.Errorf("failed to marshal permissions: %w", err)
	}

	query := `
	INSERT INTO api_keys (user_id, key_hash, admin, roles, permissions, expires_at)
	VALUES (?, ?, ?, ?, ?, ?)
	ON CONFLICT(user_id) DO UPDATE SET
		key_hash = excluded.key_hash,
		admin = excluded.admin,
		roles = excluded.roles,
		permissions = excluded.permissions,
		expires_at = excluded.expires_at,
		revoked_at = NULL
	`

	_, err = s.db.ExecContext(
		ctx,
		query,
		userID,
		keyHash,
		admin,
		rolesJSON,
		permissionsJSON,
		expiresAt,
	)
	return err
}

// RevokeAPIKey revokes an API key
func (s *DatabaseAuthStore) RevokeAPIKey(ctx context.Context, userID string) error {
	query := `UPDATE api_keys SET revoked_at = CURRENT_TIMESTAMP WHERE user_id = ?`
	_, err := s.db.ExecContext(ctx, query, userID)
	return err
}

// ListUsers returns all active users
func (s *DatabaseAuthStore) ListUsers(ctx context.Context) ([]APIKeyRecord, error) {
	query := `
	SELECT id, user_id, key_hash, admin, roles, permissions, created_at, expires_at, revoked_at
	FROM api_keys
	WHERE revoked_at IS NULL
	ORDER BY created_at DESC
	`

	rows, err := s.db.QueryContext(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("failed to query users: %w", err)
	}
	defer func() { _ = rows.Close() }()

	var users []APIKeyRecord
	for rows.Next() {
		var user APIKeyRecord
		err := rows.Scan(
			&user.ID, &user.UserID, &user.KeyHash, &user.Admin,
			&user.Roles, &user.Permissions, &user.CreatedAt,
			&user.ExpiresAt, &user.RevokedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan user: %w", err)
		}
		users = append(users, user)
	}

	if err = rows.Err(); err != nil {
		return nil, fmt.Errorf("error iterating users: %w", err)
	}

	return users, nil
}

// CleanupExpiredKeys removes expired and revoked keys
func (s *DatabaseAuthStore) CleanupExpiredKeys(ctx context.Context) error {
	query := `
	DELETE FROM api_keys 
	WHERE (revoked_at IS NOT NULL AND revoked_at < datetime('now', '-30 days'))
	   OR (expires_at IS NOT NULL AND expires_at < datetime('now', '-7 days'))
	`

	_, err := s.db.ExecContext(ctx, query)
	return err
}

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