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 import time from contextlib import contextmanager from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, Iterator, List, Tuple AUTH_DB_FILENAME = "automl_accounts.sqlite3" AUTH_DB_SNAPSHOT_FILENAME = "automl_accounts.snapshot.sqlite3" AUTH_DB_PREVIOUS_SNAPSHOT_FILENAME = "automl_accounts.snapshot.previous.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 _RUNTIME_DB_PREPARED = 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 _snapshot_database_enabled() -> bool: configured = os.getenv("AUTOML_DB_SNAPSHOT_MODE") if configured is not None: return env_truthy("AUTOML_DB_SNAPSHOT_MODE") root = storage_root() return root == Path("/data") or Path("/data") in root.parents def _persistent_storage_requires_existing_database(root: Path) -> bool: """Fail closed when authenticated snapshot storage appears temporarily empty.""" return bool( _snapshot_database_enabled() and env_truthy("AUTOML_REQUIRE_AUTH", False) and not env_truthy("AUTOML_ALLOW_EMPTY_PERSISTENT_STORAGE_INIT", False) ) def _storage_mount_wait_seconds() -> float: raw = os.getenv("AUTOML_STORAGE_MOUNT_WAIT_SECONDS", "30").strip() try: return min(120.0, max(0.0, float(raw))) except ValueError: return 30.0 def _wait_for_persistent_database(candidates: Tuple[Path, ...]) -> None: deadline = time.monotonic() + _storage_mount_wait_seconds() while not any(candidate.is_file() for candidate in candidates): remaining = deadline - time.monotonic() if remaining <= 0: break time.sleep(min(0.25, remaining)) def _runtime_auth_db_path() -> Path: if not _snapshot_database_enabled(): return auth_db_path() configured = os.getenv("AUTOML_DB_RUNTIME_DIR", "").strip() runtime_root = ( Path(configured).expanduser() if configured else Path(tempfile.gettempdir()) / "automl_gateway_state" ) root_key = hashlib.sha256(str(storage_root()).encode("utf-8")).hexdigest()[:12] return runtime_root / root_key / AUTH_DB_FILENAME def _database_integrity_ok(path: Path) -> bool: if not path.is_file(): return False conn: sqlite3.Connection | None = None try: uri = f"file:{path.resolve()}?mode=ro&immutable=1" conn = sqlite3.connect(uri, uri=True, timeout=5) row = conn.execute("PRAGMA quick_check").fetchone() return bool(row and str(row[0]).lower() == "ok") except sqlite3.DatabaseError: return False finally: if conn is not None: conn.close() def _copy_valid_database(source: Path, destination: Path) -> bool: if not source.is_file(): return False destination.parent.mkdir(parents=True, exist_ok=True) staged = destination.with_suffix(destination.suffix + ".restore") staged.unlink(missing_ok=True) try: shutil.copy2(source, staged) if not _database_integrity_ok(staged): return False os.replace(staged, destination) return True finally: staged.unlink(missing_ok=True) def _salvage_database(source: Path, destination: Path) -> bool: """Rebuild readable application tables without trusting damaged indexes.""" if not source.is_file(): return False staged = destination.with_suffix(destination.suffix + ".salvage") staged.unlink(missing_ok=True) source_conn = None destination_conn = None tables = ( ( "users", frozenset( { "user_id", "username", "username_normalized", "password_hash", "password_salt", } ), ), ( "saved_models", frozenset( { "model_id", "user_id", "manifest_json", "storage_path", "sha256", } ), ), ( "scheduled_jobs", frozenset( { "request_id", "user_id", "idempotency_key", "provider", "status", } ), ), ) try: uri = f"file:{source.resolve()}?mode=ro&immutable=1" source_conn = sqlite3.connect(uri, uri=True, timeout=5) source_conn.row_factory = sqlite3.Row source_conn.execute("PRAGMA writable_schema=ON") destination.parent.mkdir(parents=True, exist_ok=True) destination_conn = sqlite3.connect(str(staged), timeout=30) destination_conn.row_factory = sqlite3.Row destination_conn.execute("PRAGMA journal_mode=DELETE") destination_conn.execute("PRAGMA foreign_keys=ON") _create_schema(destination_conn) for table, required_columns in tables: columns = [ str(row["name"]) for row in source_conn.execute( f"PRAGMA table_info({table})" ).fetchall() ] if not required_columns.issubset(columns) or any( not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", column) for column in columns ): return False rows = source_conn.execute(f"SELECT * FROM {table}").fetchall() source_count = int( source_conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] ) if len(rows) != source_count: return False if rows: quoted_columns = ", ".join(f'"{column}"' for column in columns) placeholders = ", ".join("?" for _ in columns) destination_conn.executemany( f"INSERT INTO {table} ({quoted_columns}) VALUES ({placeholders})", [tuple(row[column] for column in columns) for row in rows], ) if destination_conn.execute("PRAGMA foreign_key_check").fetchone(): return False destination_conn.commit() destination_conn.close() destination_conn = None if not _database_integrity_ok(staged): return False os.replace(staged, destination) return True except (OSError, sqlite3.DatabaseError, TypeError, ValueError): return False finally: if source_conn is not None: source_conn.close() if destination_conn is not None: destination_conn.close() staged.unlink(missing_ok=True) def _prepare_runtime_database() -> None: global _RUNTIME_DB_PREPARED if _RUNTIME_DB_PREPARED or not _snapshot_database_enabled(): return runtime_path = _runtime_auth_db_path() runtime_path.parent.mkdir(parents=True, exist_ok=True) if _database_integrity_ok(runtime_path): _RUNTIME_DB_PREPARED = True return runtime_path.unlink(missing_ok=True) root = storage_root() candidates = ( root / AUTH_DB_SNAPSHOT_FILENAME, root / AUTH_DB_PREVIOUS_SNAPSHOT_FILENAME, auth_db_path(), ) if ( not any(candidate.is_file() for candidate in candidates) and _persistent_storage_requires_existing_database(root) ): _wait_for_persistent_database(candidates) if not any(candidate.is_file() for candidate in candidates): raise StorageConfigurationError( "Persistent account storage is unavailable or unexpectedly empty. " "The gateway stopped without creating a replacement database." ) if any(candidate.is_file() for candidate in candidates): for candidate in candidates: if _copy_valid_database(candidate, runtime_path) or _salvage_database( candidate, runtime_path, ): _RUNTIME_DB_PREPARED = True return raise StorageConfigurationError( "The persisted account database and its recovery snapshots failed " "integrity checks. The gateway stopped without replacing account data." ) _RUNTIME_DB_PREPARED = True def _persist_runtime_database(conn: sqlite3.Connection) -> None: if not _snapshot_database_enabled(): return root = storage_root() root.mkdir(parents=True, exist_ok=True) runtime_path = _runtime_auth_db_path() local_snapshot = runtime_path.with_suffix(".snapshot.tmp") durable_snapshot = root / AUTH_DB_SNAPSHOT_FILENAME previous_snapshot = root / AUTH_DB_PREVIOUS_SNAPSHOT_FILENAME staged_snapshot = root / f"{AUTH_DB_SNAPSHOT_FILENAME}.next" staged_previous = root / f"{AUTH_DB_PREVIOUS_SNAPSHOT_FILENAME}.next" local_snapshot.unlink(missing_ok=True) staged_snapshot.unlink(missing_ok=True) staged_previous.unlink(missing_ok=True) try: destination = sqlite3.connect(str(local_snapshot)) try: conn.backup(destination) finally: destination.close() if not _database_integrity_ok(local_snapshot): raise StorageConfigurationError( "The account database snapshot failed its integrity check." ) if _database_integrity_ok(durable_snapshot): shutil.copy2(durable_snapshot, staged_previous) os.replace(staged_previous, previous_snapshot) shutil.copy2(local_snapshot, staged_snapshot) if not _database_integrity_ok(staged_snapshot): raise StorageConfigurationError( "The durable account database snapshot failed verification." ) os.replace(staged_snapshot, durable_snapshot) finally: local_snapshot.unlink(missing_ok=True) staged_snapshot.unlink(missing_ok=True) staged_previous.unlink(missing_ok=True) def _utcnow() -> str: return datetime.now(timezone.utc).replace(microsecond=0).isoformat() @contextmanager def _connect() -> Iterator[sqlite3.Connection]: root = storage_root() _prepare_runtime_database() root.mkdir(parents=True, exist_ok=True) database_path = _runtime_auth_db_path() database_path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(str(database_path), timeout=30) conn.row_factory = sqlite3.Row conn.execute("PRAGMA journal_mode=DELETE") conn.execute("PRAGMA foreign_keys=ON") conn.execute("PRAGMA busy_timeout=30000") initial_changes = conn.total_changes snapshot_needs_refresh = ( _snapshot_database_enabled() and not _database_integrity_ok( storage_root() / AUTH_DB_SNAPSHOT_FILENAME ) ) try: yield conn except Exception: conn.rollback() raise else: conn.commit() if conn.total_changes > initial_changes or snapshot_needs_refresh: _persist_runtime_database(conn) finally: conn.close() 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.execute( """ CREATE TABLE IF NOT EXISTS scheduled_jobs ( request_id TEXT PRIMARY KEY, user_id TEXT NOT NULL, idempotency_key TEXT NOT NULL, profile TEXT NOT NULL, provider TEXT NOT NULL, external_job_id TEXT, status TEXT NOT NULL, request_json TEXT NOT NULL, result_json TEXT, display_name TEXT NOT NULL, result_location TEXT, remote_result_location TEXT, package_sha256 TEXT, routing_reason TEXT, cancellation_requested INTEGER NOT NULL DEFAULT 0, cleanup_pending INTEGER NOT NULL DEFAULT 0, queued_at TEXT NOT NULL, submitted_at TEXT, started_at TEXT, finished_at TEXT, cancelled_at TEXT, updated_at TEXT NOT NULL, FOREIGN KEY(user_id) REFERENCES users(user_id) ON DELETE CASCADE, UNIQUE(user_id, idempotency_key) ) """ ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_scheduled_jobs_active " "ON scheduled_jobs(status, provider, updated_at)" ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_scheduled_jobs_user_created " "ON scheduled_jobs(user_id, queued_at)" ) scheduled_columns = { row["name"] for row in conn.execute("PRAGMA table_info(scheduled_jobs)").fetchall() } scheduled_migrations = { "remote_result_location": ( "ALTER TABLE scheduled_jobs ADD COLUMN remote_result_location TEXT" ), "cleanup_pending": ( "ALTER TABLE scheduled_jobs ADD COLUMN cleanup_pending INTEGER NOT NULL DEFAULT 0" ), "routing_reason": ( "ALTER TABLE scheduled_jobs ADD COLUMN routing_reason TEXT" ), } for column_name, statement in scheduled_migrations.items(): if column_name not in scheduled_columns: conn.execute(statement) 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) _SCHEDULED_JOB_UPDATE_FIELDS = frozenset( { "provider", "external_job_id", "status", "result_json", "result_location", "remote_result_location", "package_sha256", "routing_reason", "cancellation_requested", "cleanup_pending", "submitted_at", "started_at", "finished_at", "cancelled_at", } ) _ACTIVE_SCHEDULED_STATUSES = ( "queued", "submitting", "pending", "running", "importing", "cancelling", "cleanup_pending", ) def _scheduled_job_dict(row: sqlite3.Row | None) -> Dict[str, Any] | None: if row is None: return None value = dict(row) value["cancellation_requested"] = bool( value.get("cancellation_requested", 0) ) value["cleanup_pending"] = bool(value.get("cleanup_pending", 0)) for key in ("request_json", "result_json"): raw = value.get(key) if not raw: value[key.removesuffix("_json") + "_payload"] = None continue try: parsed = json.loads(raw) except json.JSONDecodeError: parsed = None value[key.removesuffix("_json") + "_payload"] = parsed return value def create_or_get_scheduled_job( username: str, *, request_id: str, idempotency_key: str, profile: str, request_payload: Dict[str, Any], display_name: str, ) -> Dict[str, Any]: """Persist one scheduler request, returning the prior idempotent row.""" 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.") normalized_request = str(request_id or "").strip() normalized_key = str(idempotency_key or "").strip() normalized_profile = str(profile or "").strip().lower() if not normalized_request or not normalized_key: raise ValueError("Scheduler request and idempotency IDs are required.") if normalized_profile not in {"cpu", "gpu"}: raise ValueError("Scheduler profile must be cpu or gpu.") encoded_request = json.dumps( dict(request_payload), separators=(",", ":"), sort_keys=True ) now = _utcnow() with _DB_LOCK: with _connect() as conn: existing = conn.execute( """ SELECT jobs.*, users.username FROM scheduled_jobs AS jobs JOIN users ON users.user_id = jobs.user_id WHERE jobs.user_id = ? AND jobs.idempotency_key = ? """, (user["user_id"], normalized_key), ).fetchone() if existing is not None: if ( str(existing["profile"]) != normalized_profile or str(existing["request_json"]) != encoded_request ): raise ValueError( "The idempotency key is already bound to a different scheduler request." ) return _scheduled_job_dict(existing) or {} conn.execute( """ INSERT INTO scheduled_jobs ( request_id, user_id, idempotency_key, profile, provider, status, request_json, display_name, queued_at, updated_at ) VALUES (?, ?, ?, ?, 'local', 'queued', ?, ?, ?, ?) """, ( normalized_request, user["user_id"], normalized_key, normalized_profile, encoded_request, str(display_name or "AutoML model")[:120], now, now, ), ) conn.commit() return get_scheduled_job(username, normalized_request) or {} def get_scheduled_job( username: str | None, request_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 request_id: return None with _DB_LOCK: with _connect() as conn: row = conn.execute( """ SELECT jobs.*, users.username FROM scheduled_jobs AS jobs JOIN users ON users.user_id = jobs.user_id WHERE jobs.user_id = ? AND jobs.request_id = ? """, (user["user_id"], str(request_id)), ).fetchone() return _scheduled_job_dict(row) def update_scheduled_job(request_id: str, **changes: Any) -> Dict[str, Any]: """Update only the persisted provider lifecycle fields.""" initialize_storage() unknown = set(changes) - _SCHEDULED_JOB_UPDATE_FIELDS if unknown: raise ValueError(f"Unsupported scheduled job fields: {sorted(unknown)}") if not changes: raise ValueError("At least one scheduled job field must be updated.") assignments: List[str] = [] values: List[Any] = [] for key, value in changes.items(): column = key if key == "result_json" and isinstance(value, (dict, list)): value = json.dumps(value, separators=(",", ":"), sort_keys=True) if key == "cancellation_requested": value = 1 if value else 0 assignments.append(f"{column} = ?") values.append(value) assignments.append("updated_at = ?") values.extend([_utcnow(), str(request_id)]) with _DB_LOCK: with _connect() as conn: cursor = conn.execute( f"UPDATE scheduled_jobs SET {', '.join(assignments)} WHERE request_id = ?", values, ) if cursor.rowcount != 1: raise KeyError("Scheduled job was not found.") row = conn.execute( """ SELECT jobs.*, users.username FROM scheduled_jobs AS jobs JOIN users ON users.user_id = jobs.user_id WHERE jobs.request_id = ? """, (str(request_id),), ).fetchone() conn.commit() return _scheduled_job_dict(row) or {} def list_active_scheduled_jobs() -> List[Dict[str, Any]]: initialize_storage() placeholders = ",".join("?" for _ in _ACTIVE_SCHEDULED_STATUSES) with _DB_LOCK: with _connect() as conn: rows = conn.execute( f""" SELECT jobs.*, users.username FROM scheduled_jobs AS jobs JOIN users ON users.user_id = jobs.user_id WHERE jobs.status IN ({placeholders}) ORDER BY jobs.queued_at, jobs.request_id """, _ACTIVE_SCHEDULED_STATUSES, ).fetchall() return [value for row in rows if (value := _scheduled_job_dict(row))] def request_scheduled_job_cancellation( username: str, request_id: str ) -> Dict[str, Any]: record = get_scheduled_job(username, request_id) if record is None: raise PermissionError("Scheduled job was not found for this account.") return update_scheduled_job( request_id, cancellation_requested=True, status=( "cancelling" if record.get("status") in {"pending", "running", "submitting"} else record.get("status") ), ) 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()