auth_storage.py

from __future__ import annotations

import base64
import binascii
import hashlib
import hmac
import json
import os
import re
import secrets
import shutil
import sqlite3
import tempfile
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Tuple


AUTH_DB_FILENAME = "automl_accounts.sqlite3"
DEFAULT_MODEL_CHUNK_SIZE = 4 * 1024 * 1024
ENCRYPTED_MODEL_MAGIC = b"AUTOMLENC1\n"
USERNAME_RE = re.compile(r"^[A-Za-z0-9_.@-]{3,64}$")
ACCOUNT_STATUS_PENDING = "pending"
ACCOUNT_STATUS_APPROVED = "approved"
ACCOUNT_STATUS_DENIED = "denied"
ACCOUNT_STATUSES = frozenset(
    {ACCOUNT_STATUS_PENDING, ACCOUNT_STATUS_APPROVED, ACCOUNT_STATUS_DENIED}
)
ADMIN_USERNAMES = frozenset({"zukhriddin", "alex", "vishnu", "snigdha", "flex"})

_DB_LOCK = threading.RLock()
_BOOTSTRAPPED = False
_PROCESS_SESSION_SECRET = secrets.token_bytes(32)
_BROWSER_STATE_CONTEXT = b"automl-browser-state-v1"


class StorageConfigurationError(RuntimeError):
    """Raised when secure account/model storage is not configured correctly."""


def env_truthy(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def storage_root() -> Path:
    configured = os.getenv("AUTOML_STORAGE_DIR", "").strip()
    if configured:
        return Path(configured).expanduser()
    data_root = Path("/data")
    if data_root.is_dir() and os.access(data_root, os.W_OK):
        return data_root / "automl"
    return Path.cwd() / ".automl_storage"


def auth_db_path() -> Path:
    return storage_root() / AUTH_DB_FILENAME


def _utcnow() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def _connect() -> sqlite3.Connection:
    root = storage_root()
    root.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(auth_db_path()), timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    conn.execute("PRAGMA busy_timeout=30000")
    return conn


def _create_schema(conn: sqlite3.Connection) -> None:
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS users (
            user_id TEXT PRIMARY KEY,
            username TEXT NOT NULL,
            username_normalized TEXT NOT NULL UNIQUE,
            password_alg TEXT NOT NULL,
            password_hash TEXT NOT NULL,
            password_salt TEXT NOT NULL,
            password_iterations INTEGER NOT NULL,
            session_version INTEGER NOT NULL DEFAULT 1,
            account_status TEXT NOT NULL DEFAULT 'approved',
            requested_at TEXT,
            reviewed_at TEXT,
            reviewed_by TEXT,
            created_at TEXT NOT NULL,
            last_login_at TEXT
        )
        """
    )
    user_columns = {str(row[1]) for row in conn.execute("PRAGMA table_info(users)")}
    migrations = {
        "session_version": (
            "ALTER TABLE users ADD COLUMN session_version INTEGER NOT NULL DEFAULT 1"
        ),
        "account_status": (
            "ALTER TABLE users ADD COLUMN account_status TEXT NOT NULL DEFAULT 'approved'"
        ),
        "requested_at": "ALTER TABLE users ADD COLUMN requested_at TEXT",
        "reviewed_at": "ALTER TABLE users ADD COLUMN reviewed_at TEXT",
        "reviewed_by": "ALTER TABLE users ADD COLUMN reviewed_by TEXT",
    }
    for column_name, statement in migrations.items():
        if column_name not in user_columns:
            conn.execute(statement)
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS saved_models (
            model_id TEXT PRIMARY KEY,
            user_id TEXT NOT NULL,
            display_name TEXT NOT NULL,
            artifact_type TEXT,
            predictor TEXT,
            task_type TEXT,
            target_column TEXT,
            manifest_json TEXT NOT NULL,
            storage_path TEXT NOT NULL,
            storage_encrypted INTEGER NOT NULL,
            storage_key_id TEXT,
            source_package_name TEXT,
            byte_size INTEGER NOT NULL,
            sha256 TEXT NOT NULL,
            created_at TEXT NOT NULL,
            last_used_at TEXT,
            FOREIGN KEY(user_id) REFERENCES users(user_id) ON DELETE CASCADE
        )
        """
    )
    conn.execute("CREATE INDEX IF NOT EXISTS idx_saved_models_user_created ON saved_models(user_id, created_at)")
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_users_account_status "
        "ON users(account_status, requested_at)"
    )
    conn.commit()


def normalize_username(username: str | None) -> str:
    return str(username or "").strip()


def _username_key(username: str | None) -> str:
    return normalize_username(username).lower()


def is_admin(username: str | None) -> bool:
    return _username_key(username) in ADMIN_USERNAMES


def _is_approved_user(user: Dict[str, Any] | sqlite3.Row | None) -> bool:
    return bool(
        user
        and str(user["account_status"] or ACCOUNT_STATUS_APPROVED).lower()
        == ACCOUNT_STATUS_APPROVED
    )


def validate_username(username: str) -> None:
    if not USERNAME_RE.fullmatch(username):
        raise ValueError("Usernames must be 3-64 characters using letters, numbers, dots, underscores, hyphens, or @.")


def min_password_length() -> int:
    raw = os.getenv("AUTOML_MIN_PASSWORD_LENGTH", "10").strip()
    try:
        return max(8, int(raw))
    except ValueError:
        return 10


def validate_password(password: str) -> None:
    if len(password or "") < min_password_length():
        raise ValueError(f"Passwords must be at least {min_password_length()} characters.")


def _password_iterations() -> int:
    raw = os.getenv("AUTOML_PASSWORD_HASH_ITERATIONS", "390000").strip()
    try:
        return max(200000, int(raw))
    except ValueError:
        return 390000


def _hash_password(password: str, salt: bytes | None = None, iterations: int | None = None) -> Tuple[str, str, int]:
    salt = salt or secrets.token_bytes(16)
    iterations = iterations or _password_iterations()
    digest = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, iterations)
    return (
        base64.b64encode(digest).decode("ascii"),
        base64.b64encode(salt).decode("ascii"),
        iterations,
    )


def _verify_password(password: str, stored_hash: str, stored_salt: str, iterations: int) -> bool:
    try:
        salt = base64.b64decode(stored_salt.encode("ascii"), validate=True)
        expected = base64.b64decode(stored_hash.encode("ascii"), validate=True)
    except (binascii.Error, ValueError):
        return False
    digest = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, int(iterations))
    return hmac.compare_digest(digest, expected)


def initialize_storage() -> None:
    with _DB_LOCK:
        with _connect() as conn:
            _create_schema(conn)
        _bootstrap_users_once()


def _bootstrap_users_once() -> None:
    global _BOOTSTRAPPED
    if _BOOTSTRAPPED:
        return
    payload = os.getenv("AUTOML_BOOTSTRAP_USERS_JSON", "").strip()
    if not payload:
        _BOOTSTRAPPED = True
        return
    try:
        loaded = json.loads(payload)
    except json.JSONDecodeError as exc:
        raise StorageConfigurationError("AUTOML_BOOTSTRAP_USERS_JSON must be valid JSON.") from exc

    if isinstance(loaded, dict):
        entries = [{"username": username, "password": password} for username, password in loaded.items()]
    elif isinstance(loaded, list):
        entries = loaded
    else:
        raise StorageConfigurationError("AUTOML_BOOTSTRAP_USERS_JSON must be a JSON object or list.")

    for entry in entries:
        if not isinstance(entry, dict):
            raise StorageConfigurationError("Each bootstrap user must be an object with username and password.")
        username = normalize_username(str(entry.get("username", "")))
        password = str(entry.get("password", ""))
        create_user(username, password, if_missing=True)
    _BOOTSTRAPPED = True


def get_user(username: str | None) -> Dict[str, Any] | None:
    initialize_storage()
    key = _username_key(username)
    if not key:
        return None
    with _DB_LOCK:
        with _connect() as conn:
            row = conn.execute(
                "SELECT * FROM users WHERE username_normalized = ?",
                (key,),
            ).fetchone()
    return dict(row) if row else None


def get_user_by_id(user_id: str | None) -> Dict[str, Any] | None:
    initialize_storage()
    normalized = str(user_id or "").strip()
    if not normalized:
        return None
    with _DB_LOCK:
        with _connect() as conn:
            row = conn.execute(
                "SELECT * FROM users WHERE user_id = ?",
                (normalized,),
            ).fetchone()
    return dict(row) if row else None


def user_count() -> int:
    initialize_storage()
    with _DB_LOCK:
        with _connect() as conn:
            return int(conn.execute("SELECT COUNT(*) FROM users").fetchone()[0])


def create_user(username: str, password: str, *, if_missing: bool = False) -> Dict[str, Any]:
    """Create an approved account for trusted bootstrap/internal provisioning."""

    username = normalize_username(username)
    validate_username(username)
    validate_password(password)
    key = _username_key(username)
    with _DB_LOCK:
        with _connect() as conn:
            _create_schema(conn)
            existing = conn.execute(
                "SELECT * FROM users WHERE username_normalized = ?",
                (key,),
            ).fetchone()
            if existing:
                if if_missing:
                    return dict(existing)
                raise ValueError("That username already exists.")
            password_hash, password_salt, iterations = _hash_password(password)
            now = _utcnow()
            user_id = secrets.token_hex(16)
            conn.execute(
                """
                INSERT INTO users (
                    user_id, username, username_normalized, password_alg, password_hash,
                    password_salt, password_iterations, account_status, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    user_id,
                    username,
                    key,
                    "pbkdf2_sha256",
                    password_hash,
                    password_salt,
                    iterations,
                    ACCOUNT_STATUS_APPROVED,
                    now,
                ),
            )
            conn.commit()
            return dict(
                conn.execute(
                    "SELECT * FROM users WHERE username_normalized = ?",
                    (key,),
                ).fetchone()
            )


def request_account_creation(username: str, password: str) -> Dict[str, Any]:
    """Store a self-service signup as pending administrator review."""

    initialize_storage()
    username = normalize_username(username)
    validate_username(username)
    validate_password(password)
    if is_admin(username):
        raise ValueError(
            "That username is reserved for an administrator and must be provisioned by the deployment owner."
        )
    key = _username_key(username)
    with _DB_LOCK:
        with _connect() as conn:
            existing = conn.execute(
                "SELECT account_status FROM users WHERE username_normalized = ?",
                (key,),
            ).fetchone()
            if existing:
                status = str(existing["account_status"] or ACCOUNT_STATUS_APPROVED)
                if status == ACCOUNT_STATUS_PENDING:
                    raise ValueError(
                        "An account request for that username is already awaiting administrator approval."
                    )
                if status == ACCOUNT_STATUS_DENIED:
                    raise ValueError(
                        "That request was denied. Log in with those credentials to resend or remove it."
                    )
                raise ValueError("That username already exists.")

            password_hash, password_salt, iterations = _hash_password(password)
            now = _utcnow()
            user_id = secrets.token_hex(16)
            conn.execute(
                """
                INSERT INTO users (
                    user_id, username, username_normalized, password_alg, password_hash,
                    password_salt, password_iterations, session_version, account_status,
                    requested_at, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    user_id,
                    username,
                    key,
                    "pbkdf2_sha256",
                    password_hash,
                    password_salt,
                    iterations,
                    1,
                    ACCOUNT_STATUS_PENDING,
                    now,
                    now,
                ),
            )
            conn.commit()
            return dict(
                conn.execute(
                    "SELECT * FROM users WHERE username_normalized = ?",
                    (key,),
                ).fetchone()
            )


def account_authentication_status(username: str, password: str) -> str:
    """Return an account status only after the supplied password verifies."""

    initialize_storage()
    username = normalize_username(username)
    if not username or not password:
        return "invalid"
    key = _username_key(username)
    with _DB_LOCK:
        with _connect() as conn:
            row = conn.execute(
                "SELECT * FROM users WHERE username_normalized = ?",
                (key,),
            ).fetchone()
            if row is None or not _verify_password(
                password,
                str(row["password_hash"]),
                str(row["password_salt"]),
                int(row["password_iterations"]),
            ):
                return "invalid"
            status = str(row["account_status"] or ACCOUNT_STATUS_APPROVED).lower()
            if status not in ACCOUNT_STATUSES:
                return "invalid"
            if status == ACCOUNT_STATUS_APPROVED:
                conn.execute(
                    "UPDATE users SET last_login_at = ? WHERE user_id = ?",
                    (_utcnow(), row["user_id"]),
                )
                conn.commit()
            return status


def list_pending_account_requests() -> List[Dict[str, Any]]:
    initialize_storage()
    with _DB_LOCK:
        with _connect() as conn:
            rows = conn.execute(
                """
                SELECT username, requested_at
                FROM users
                WHERE account_status = ?
                ORDER BY requested_at ASC, username_normalized ASC
                """,
                (ACCOUNT_STATUS_PENDING,),
            ).fetchall()
    return [dict(row) for row in rows]


def review_account_request(
    admin_username: str,
    username: str,
    decision: str,
) -> Dict[str, Any]:
    initialize_storage()
    admin = get_user(admin_username)
    if admin is None or not is_admin(admin_username) or not _is_approved_user(admin):
        raise PermissionError(
            "Only an approved administrator can review account requests."
        )
    decision = str(decision or "").strip().lower()
    if decision not in {ACCOUNT_STATUS_APPROVED, ACCOUNT_STATUS_DENIED}:
        raise ValueError("Decision must be approved or denied.")
    key = _username_key(username)
    if not key:
        raise ValueError("Select a pending account request.")
    if is_admin(username):
        raise ValueError(
            "Administrator accounts must be provisioned through the bootstrap secret."
        )

    with _DB_LOCK:
        with _connect() as conn:
            row = conn.execute(
                "SELECT * FROM users WHERE username_normalized = ?",
                (key,),
            ).fetchone()
            if row is None or str(row["account_status"]) != ACCOUNT_STATUS_PENDING:
                raise ValueError("That account request is no longer pending.")
            cursor = conn.execute(
                """
                UPDATE users
                SET account_status = ?, reviewed_at = ?, reviewed_by = ?
                WHERE user_id = ? AND account_status = ?
                """,
                (
                    decision,
                    _utcnow(),
                    normalize_username(admin["username"]),
                    row["user_id"],
                    ACCOUNT_STATUS_PENDING,
                ),
            )
            if cursor.rowcount != 1:
                raise ValueError("That account request is no longer pending.")
            conn.commit()
            return dict(
                conn.execute(
                    "SELECT * FROM users WHERE user_id = ?",
                    (row["user_id"],),
                ).fetchone()
            )


def resubmit_denied_account(username: str, password: str) -> None:
    if account_authentication_status(username, password) != ACCOUNT_STATUS_DENIED:
        raise PermissionError("The denied account credentials could not be verified.")
    key = _username_key(username)
    with _DB_LOCK:
        with _connect() as conn:
            cursor = conn.execute(
                """
                UPDATE users
                SET account_status = ?, requested_at = ?, reviewed_at = NULL,
                    reviewed_by = NULL
                WHERE username_normalized = ? AND account_status = ?
                """,
                (ACCOUNT_STATUS_PENDING, _utcnow(), key, ACCOUNT_STATUS_DENIED),
            )
            if cursor.rowcount != 1:
                raise ValueError("That account is no longer denied.")
            conn.commit()


def delete_denied_account(username: str, password: str) -> None:
    if account_authentication_status(username, password) != ACCOUNT_STATUS_DENIED:
        raise PermissionError("The denied account credentials could not be verified.")
    key = _username_key(username)
    with _DB_LOCK:
        with _connect() as conn:
            row = conn.execute(
                "SELECT user_id FROM users "
                "WHERE username_normalized = ? AND account_status = ?",
                (key, ACCOUNT_STATUS_DENIED),
            ).fetchone()
            if row is None:
                raise ValueError("That account is no longer denied.")
            user_id = str(row["user_id"])
            conn.execute("DELETE FROM users WHERE user_id = ?", (user_id,))
            conn.commit()
    model_dir = _model_dir_for_user(user_id)
    if model_dir.is_dir():
        shutil.rmtree(model_dir)


def signups_allowed() -> bool:
    return env_truthy("AUTOML_ALLOW_SIGNUPS", False)


def app_auth_required() -> bool:
    if env_truthy("AUTOML_AUTH_DISABLED", False):
        return False
    initialize_storage()
    return env_truthy("AUTOML_REQUIRE_AUTH", True) or signups_allowed() or user_count() > 0


def authenticate_user(username: str, password: str, *, create_if_missing: bool = False) -> bool:
    # Retained for API compatibility. Self-service accounts must be submitted
    # explicitly and approved before authentication can succeed.
    del create_if_missing
    return (
        account_authentication_status(username, password)
        == ACCOUNT_STATUS_APPROVED
    )


def update_password(username: str, old_password: str, new_password: str) -> None:
    initialize_storage()
    username = normalize_username(username)
    if not username:
        raise PermissionError("Sign in before changing your password.")
    validate_password(new_password)
    key = _username_key(username)
    with _DB_LOCK:
        with _connect() as conn:
            row = conn.execute(
                "SELECT * FROM users WHERE username_normalized = ?",
                (key,),
            ).fetchone()
            if row is None:
                raise PermissionError("The signed-in user was not found.")
            if not _is_approved_user(row):
                raise PermissionError("Only approved accounts can change passwords.")
            if not _verify_password(
                old_password,
                str(row["password_hash"]),
                str(row["password_salt"]),
                int(row["password_iterations"]),
            ):
                raise PermissionError("The current password is incorrect.")
            password_hash, password_salt, iterations = _hash_password(new_password)
            conn.execute(
                """
                UPDATE users
                SET password_hash = ?, password_salt = ?, password_iterations = ?, password_alg = ?,
                    session_version = session_version + 1
                WHERE user_id = ?
                """,
                (
                    password_hash,
                    password_salt,
                    iterations,
                    "pbkdf2_sha256",
                    row["user_id"],
                ),
            )
            conn.commit()


def get_gradio_auth():
    if env_truthy("AUTOML_AUTH_DISABLED", False):
        return None
    initialize_storage()
    if user_count() == 0 and not signups_allowed():
        if env_truthy("AUTOML_REQUIRE_AUTH", False):
            raise StorageConfigurationError(
                "No AutoML users are configured. Set AUTOML_BOOTSTRAP_USERS_JSON as a Hugging Face Space Secret."
            )
        return None
    return lambda username, password: authenticate_user(
        username,
        password,
        create_if_missing=False,
    )


def _session_secret() -> bytes:
    configured = os.getenv("AUTOML_SESSION_SECRET", "").strip()
    if configured:
        padded = configured + ("=" * ((4 - len(configured) % 4) % 4))
        try:
            secret = base64.urlsafe_b64decode(padded.encode("ascii"))
        except (binascii.Error, ValueError) as exc:
            raise StorageConfigurationError(
                "AUTOML_SESSION_SECRET must be URL-safe base64 for at least 32 random bytes."
            ) from exc
        if len(secret) < 32:
            raise StorageConfigurationError("AUTOML_SESSION_SECRET must decode to at least 32 bytes.")
        return secret
    model_key = _decode_model_storage_key()
    if model_key is not None:
        return hmac.new(model_key, b"automl-session-v1", hashlib.sha256).digest()
    return _PROCESS_SESSION_SECRET


def _b64url_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")


def browser_state_secret() -> str:
    """Return a stable purpose-separated key for encrypted browser state."""

    derived = hmac.new(
        _session_secret(),
        _BROWSER_STATE_CONTEXT,
        hashlib.sha256,
    ).digest()
    return _b64url_encode(derived)


def _b64url_decode(value: str) -> bytes:
    padded = value + ("=" * ((4 - len(value) % 4) % 4))
    return base64.urlsafe_b64decode(padded.encode("ascii"))


def issue_session_token(username: str) -> str:
    user = get_user(username)
    if user is None or not _is_approved_user(user):
        raise PermissionError("Cannot create a session for an unapproved user.")
    payload = {
        "user_id": user["user_id"],
        "username": user["username"],
        "session_version": int(user.get("session_version") or 1),
        "iat": int(datetime.now(timezone.utc).timestamp()),
        "nonce": secrets.token_urlsafe(16),
    }
    payload_bytes = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")
    encoded_payload = _b64url_encode(payload_bytes)
    signature = hmac.new(_session_secret(), encoded_payload.encode("ascii"), hashlib.sha256).digest()
    return f"{encoded_payload}.{_b64url_encode(signature)}"


def username_from_session_token(session_token: str | None) -> str:
    token = str(session_token or "").strip()
    if not token or "." not in token:
        return ""
    encoded_payload, encoded_signature = token.rsplit(".", 1)
    try:
        expected = hmac.new(_session_secret(), encoded_payload.encode("ascii"), hashlib.sha256).digest()
        provided = _b64url_decode(encoded_signature)
    except (binascii.Error, ValueError):
        return ""
    if not hmac.compare_digest(expected, provided):
        return ""
    try:
        payload = json.loads(_b64url_decode(encoded_payload).decode("utf-8"))
    except (json.JSONDecodeError, UnicodeDecodeError, binascii.Error, ValueError):
        return ""
    issued_at = int(payload.get("iat") or 0)
    max_age_raw = os.getenv("AUTOML_SESSION_MAX_AGE_SECONDS", str(30 * 24 * 60 * 60)).strip()
    try:
        max_age = max(3600, int(max_age_raw))
    except ValueError:
        max_age = 30 * 24 * 60 * 60
    now = int(datetime.now(timezone.utc).timestamp())
    if issued_at <= 0 or issued_at > now + 300 or now - issued_at > max_age:
        return ""
    user = get_user_by_id(payload.get("user_id"))
    if user is None or not _is_approved_user(user):
        return ""
    try:
        token_version = int(payload.get("session_version"))
        stored_version = int(user.get("session_version") or 1)
    except (TypeError, ValueError):
        return ""
    if token_version != stored_version:
        return ""
    token_username = normalize_username(payload.get("username"))
    if _username_key(token_username) != _username_key(user.get("username")):
        return ""
    return str(user["username"])


def current_username(request: Any | None = None, session_token: str | None = None) -> str:
    request_username = normalize_username(getattr(request, "username", None))
    if request_username:
        return request_username
    return username_from_session_token(session_token)


def require_authenticated_username(request: Any | None = None, session_token: str | None = None) -> str:
    username = current_username(request, session_token)
    if not username:
        raise PermissionError("Sign in before using saved cloud models.")
    return username


def _decode_model_storage_key() -> bytes | None:
    raw = os.getenv("AUTOML_MODEL_STORAGE_KEY", "").strip()
    if not raw:
        return None
    padded = raw + ("=" * ((4 - len(raw) % 4) % 4))
    try:
        key = base64.urlsafe_b64decode(padded.encode("ascii"))
    except (binascii.Error, ValueError) as exc:
        raise StorageConfigurationError(
            "AUTOML_MODEL_STORAGE_KEY must be URL-safe base64 for 32 random bytes."
        ) from exc
    if len(key) != 32:
        raise StorageConfigurationError("AUTOML_MODEL_STORAGE_KEY must decode to exactly 32 bytes.")
    return key


def encryption_enabled() -> bool:
    return _decode_model_storage_key() is not None


def _model_storage_key_id() -> str:
    key = _decode_model_storage_key()
    if key is None:
        return ""
    return hashlib.sha256(key).hexdigest()[:16]


def _chunk_aad(header: bytes, index: int) -> bytes:
    return ENCRYPTED_MODEL_MAGIC + header + index.to_bytes(8, "big")


def _encrypt_file(source: Path, destination: Path, *, model_id: str) -> Tuple[int, str]:
    key = _decode_model_storage_key()
    if key is None:
        raise StorageConfigurationError("Set AUTOML_MODEL_STORAGE_KEY before saving encrypted models.")
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    except ImportError as exc:
        raise StorageConfigurationError("Install the cryptography package to save encrypted models.") from exc

    destination.parent.mkdir(parents=True, exist_ok=True)
    nonce_prefix = secrets.token_bytes(8)
    header_obj = {
        "alg": "AES-256-GCM-CHUNKED",
        "chunk_size": DEFAULT_MODEL_CHUNK_SIZE,
        "nonce_prefix": base64.b64encode(nonce_prefix).decode("ascii"),
        "key_id": _model_storage_key_id(),
        "model_id": model_id,
    }
    header = json.dumps(header_obj, separators=(",", ":")).encode("utf-8")
    aesgcm = AESGCM(key)
    sha = hashlib.sha256()
    total = 0
    index = 0

    with source.open("rb") as src, destination.open("wb") as dst:
        dst.write(ENCRYPTED_MODEL_MAGIC)
        dst.write(len(header).to_bytes(4, "big"))
        dst.write(header)
        while True:
            chunk = src.read(DEFAULT_MODEL_CHUNK_SIZE)
            if not chunk:
                break
            sha.update(chunk)
            total += len(chunk)
            nonce = nonce_prefix + index.to_bytes(4, "big")
            encrypted = aesgcm.encrypt(nonce, chunk, _chunk_aad(header, index))
            dst.write(len(encrypted).to_bytes(4, "big"))
            dst.write(encrypted)
            index += 1
    return total, sha.hexdigest()


def _decrypt_file(source: Path, destination: Path, *, expected_size: int, expected_sha256: str) -> None:
    key = _decode_model_storage_key()
    if key is None:
        raise StorageConfigurationError("Set AUTOML_MODEL_STORAGE_KEY before loading encrypted saved models.")
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    except ImportError as exc:
        raise StorageConfigurationError("Install the cryptography package to load encrypted models.") from exc

    aesgcm = AESGCM(key)
    sha = hashlib.sha256()
    total = 0
    index = 0
    destination.parent.mkdir(parents=True, exist_ok=True)

    try:
        with source.open("rb") as src, destination.open("wb") as dst:
            magic = src.read(len(ENCRYPTED_MODEL_MAGIC))
            if magic != ENCRYPTED_MODEL_MAGIC:
                raise StorageConfigurationError("Saved model file is not in the expected encrypted format.")
            header_len_raw = src.read(4)
            if len(header_len_raw) != 4:
                raise StorageConfigurationError("Saved model encryption header is incomplete.")
            header_len = int.from_bytes(header_len_raw, "big")
            header = src.read(header_len)
            if len(header) != header_len:
                raise StorageConfigurationError("Saved model encryption header is truncated.")
            header_obj = json.loads(header.decode("utf-8"))
            nonce_prefix = base64.b64decode(header_obj["nonce_prefix"].encode("ascii"), validate=True)
            while True:
                length_raw = src.read(4)
                if not length_raw:
                    break
                if len(length_raw) != 4:
                    raise StorageConfigurationError("Saved model ciphertext length is incomplete.")
                encrypted_len = int.from_bytes(length_raw, "big")
                encrypted = src.read(encrypted_len)
                if len(encrypted) != encrypted_len:
                    raise StorageConfigurationError("Saved model ciphertext is truncated.")
                nonce = nonce_prefix + index.to_bytes(4, "big")
                chunk = aesgcm.decrypt(nonce, encrypted, _chunk_aad(header, index))
                dst.write(chunk)
                sha.update(chunk)
                total += len(chunk)
                index += 1
    except Exception:
        if destination.exists():
            destination.unlink(missing_ok=True)
        raise

    if total != int(expected_size) or sha.hexdigest() != expected_sha256:
        destination.unlink(missing_ok=True)
        raise StorageConfigurationError("Saved model integrity check failed after decryption.")


def _copy_plaintext_file(source: Path, destination: Path) -> Tuple[int, str]:
    destination.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(source, destination)
    sha = hashlib.sha256()
    total = 0
    with source.open("rb") as src:
        while True:
            chunk = src.read(DEFAULT_MODEL_CHUNK_SIZE)
            if not chunk:
                break
            sha.update(chunk)
            total += len(chunk)
    return total, sha.hexdigest()


def _model_dir_for_user(user_id: str) -> Path:
    return storage_root() / "models" / user_id


def save_model_for_user(
    username: str,
    package_path: str,
    *,
    display_name: str,
    manifest: Dict[str, Any],
    artifact_type: str,
) -> Dict[str, Any]:
    initialize_storage()
    user = get_user(username)
    if user is None or not _is_approved_user(user):
        raise PermissionError("The signed-in user was not found in the account database.")
    source = Path(package_path)
    if not source.is_file():
        raise FileNotFoundError(f"Model package was not found: {package_path}")

    model_id = secrets.token_urlsafe(12)
    encrypted = encryption_enabled()
    if encrypted:
        destination = _model_dir_for_user(user["user_id"]) / f"{model_id}.zip.enc"
        byte_size, digest = _encrypt_file(source, destination, model_id=model_id)
        key_id = _model_storage_key_id()
    elif env_truthy("AUTOML_ALLOW_UNENCRYPTED_MODEL_STORAGE", False):
        destination = _model_dir_for_user(user["user_id"]) / f"{model_id}.zip"
        byte_size, digest = _copy_plaintext_file(source, destination)
        key_id = ""
    else:
        raise StorageConfigurationError(
            "Encrypted model storage is not configured. Set AUTOML_MODEL_STORAGE_KEY as a Space Secret."
        )

    now = _utcnow()
    manifest_json = json.dumps(manifest or {}, sort_keys=True)
    with _DB_LOCK:
        with _connect() as conn:
            conn.execute(
                """
                INSERT INTO saved_models (
                    model_id, user_id, display_name, artifact_type, predictor, task_type, target_column,
                    manifest_json, storage_path, storage_encrypted, storage_key_id, source_package_name,
                    byte_size, sha256, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    model_id,
                    user["user_id"],
                    display_name[:120] or "AutoML model",
                    artifact_type or "",
                    str(manifest.get("predictor") or ""),
                    str(manifest.get("task_type") or ""),
                    str(manifest.get("target_column") or manifest.get("label_column") or ""),
                    manifest_json,
                    str(destination),
                    1 if encrypted else 0,
                    key_id,
                    source.name,
                    int(byte_size),
                    digest,
                    now,
                ),
            )
            conn.commit()
    return get_saved_model(username, model_id) or {"model_id": model_id, "display_name": display_name}


def list_saved_models(username: str | None) -> List[Dict[str, Any]]:
    initialize_storage()
    user = get_user(username)
    if user is None or not _is_approved_user(user):
        return []
    with _DB_LOCK:
        with _connect() as conn:
            rows = conn.execute(
                """
                SELECT * FROM saved_models
                WHERE user_id = ?
                ORDER BY created_at DESC, display_name ASC
                """,
                (user["user_id"],),
            ).fetchall()
    return [dict(row) for row in rows]


def get_saved_model(username: str | None, model_id: str | None) -> Dict[str, Any] | None:
    initialize_storage()
    user = get_user(username)
    if user is None or not _is_approved_user(user) or not model_id:
        return None
    with _DB_LOCK:
        with _connect() as conn:
            row = conn.execute(
                """
                SELECT * FROM saved_models
                WHERE user_id = ? AND model_id = ?
                """,
                (user["user_id"], str(model_id)),
            ).fetchone()
    return dict(row) if row else None


def manifest_from_record(record: Dict[str, Any]) -> Dict[str, Any]:
    try:
        loaded = json.loads(record.get("manifest_json") or "{}")
        return loaded if isinstance(loaded, dict) else {}
    except json.JSONDecodeError:
        return {}


def materialize_saved_model(username: str, model_id: str) -> Tuple[str, str]:
    record = get_saved_model(username, model_id)
    if record is None:
        raise PermissionError("Saved model not found for this account.")
    source = Path(record["storage_path"])
    if not source.is_file():
        raise FileNotFoundError("The saved model file is missing from storage.")

    temp_dir = Path(tempfile.mkdtemp(prefix="automl_saved_model_"))
    destination = temp_dir / f"{model_id}.zip"
    if int(record.get("storage_encrypted") or 0):
        _decrypt_file(
            source,
            destination,
            expected_size=int(record["byte_size"]),
            expected_sha256=str(record["sha256"]),
        )
    else:
        if not env_truthy("AUTOML_ALLOW_UNENCRYPTED_MODEL_STORAGE", False):
            raise StorageConfigurationError("Plaintext saved model storage is disabled.")
        shutil.copy2(source, destination)

    with _DB_LOCK:
        with _connect() as conn:
            conn.execute(
                "UPDATE saved_models SET last_used_at = ? WHERE model_id = ?",
                (_utcnow(), model_id),
            )
            conn.commit()
    return str(destination), str(temp_dir)


def delete_saved_model(username: str, model_id: str) -> None:
    record = get_saved_model(username, model_id)
    if record is None:
        raise PermissionError("Saved model not found for this account.")
    storage_path = Path(record["storage_path"])
    with _DB_LOCK:
        with _connect() as conn:
            conn.execute("DELETE FROM saved_models WHERE model_id = ?", (model_id,))
            conn.commit()
    if storage_path.is_file():
        storage_path.unlink()