AutoML / auth_storage.py
zukhriddinai's picture
Deploy public authenticated AutoML gateway
660c411 verified
Raw History Blame Contribute Delete
57.5 kB
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()