unsloth/studio/backend/auth/storage.py
Leo Borcherding 91a0df9514
Studio: make the Cloudflare tunnel opt-in (off by default) (#7046)
* Studio: make the Cloudflare tunnel opt-in (off by default)

A wildcard bind (`-H 0.0.0.0`) auto-started a public trycloudflare.com
tunnel, so exposing Studio on the LAN also published it to the public
internet. Flip the default so the tunnel is opt-in.

- `--cloudflare` is now tri-state (Optional[bool], default None = off),
  mirroring the existing --enable-tools/--disable-tools handling. Pass
  --cloudflare to expose a public HTTPS link for a wildcard bind; --secure
  still implies the tunnel.
- --secure + --no-cloudflare is still rejected as a contradiction.
- Update the parent-command guard, re-exec forwarding, startup-banner
  wording, the colab comment, README, and tests.

* Studio: update installer/setup launch hints for opt-in Cloudflare

The post-install launch hints only mentioned --secure for a public link.
Now that the tunnel is opt-in, clarify that -H 0.0.0.0 exposes the raw
port on the LAN (not a public URL), and surface --cloudflare as the
explicit opt-in for a public HTTPS link (--secure keeps the raw port
private). Applied to install.ps1, install.sh, and studio/setup.sh.

* Studio: address review - keep cloudflare tri-state + harden run re-exec

Two review points from the bots:

- Gemini: keep `cloudflare` as Optional[bool] in run_server instead of
  casting None -> False, so the startup banner can distinguish "OFF (default)"
  (unset) from "OFF (--no-cloudflare)" (explicit). `_cloudflare_flag` and the
  banner branch now carry the tri-state.
- Codex (P1): `unsloth studio run` re-execs the studio venv's console script,
  which can be an older build whose --cloudflare defaulted on; omitting the
  flag let it re-enable the tunnel. That path now forwards the default polarity
  explicitly (--no-cloudflare, or nothing under --secure since --secure implies
  the tunnel). The plain `unsloth studio` path runs the same-version in-tree
  run.py (resolved via _find_run_py), so it keeps forwarding only an explicit
  polarity and still shows the accurate "(default)" banner.

Tests updated for the tri-state banner labels, the None gate cases, and the
new re-exec forwarding.

* Studio: forward --no-cloudflare on plain re-exec too (mixed install)

Codex follow-up: _find_run_py falls back to STUDIO_HOME/.../studio/backend/
run.py when the package copy is absent, so the plain `unsloth studio` re-exec
can land on an older run.py whose --cloudflare defaults on. Forward the default
polarity explicitly there too (--no-cloudflare, or nothing under --secure),
matching the run subcommand. The common in-venv launch skips the re-exec and
still shows the tri-state "(default)" banner.

* Studio: fix launch hint - --cloudflare needs the wildcard bind

Codex P3: the launch hint listed --cloudflare next to the loopback
`unsloth studio -p 8888` command, but the tunnel only starts for wildcard
binds, so `--cloudflare` alone on 127.0.0.1 does nothing. Show
`-H 0.0.0.0 --cloudflare` in the hints (install.ps1, install.sh,
studio/setup.sh) and clarify the same in the README.

* Studio: cross-platform masked terminal password prompt helper

Per-keystroke '*' echo (POSIX termios cbreak / Windows msvcrt.getwch),
backspace editing, Ctrl-C abort, EOF handling, confirmation loop with
re-prompt on mismatch or policy failure. Pure should_prompt gate for the
--secure/--cloudflare exposure paths.

* Studio CLI: force a terminal password change before public tunnel exposure

When a launch will start the Cloudflare tunnel (--secure, or --cloudflare on
a non-api-only wildcard bind) and the admin account still has its seeded
bootstrap password, prompt for a new password in the terminal (masked with
'*', confirmed, re-prompting until valid) before any re-exec or server
exists. The change is committed in the parent so it never crosses argv or
the environment and older studio-venv children see it immediately. Without
a terminal, warn and fall back to the backend bootstrap shutdown timer.
Mirrors backend update_password semantics in one transaction: rehash,
rotate the JWT secret, clear must_change_password, revoke refresh tokens,
drop the desktop secret, then remove the stale credential files.

* Studio: terminal password gate before the public tunnel (backend backstop)

Never publish a trycloudflare URL while the seeded admin password is
active: run_server now runs a terminal password-change gate after the
tunnel decision and strictly before start_studio_tunnel. Interactive
refusal fails closed (shutdown + exit 1, mirroring the secure gate);
without a tty it warns and keeps the bootstrap deadline. Success applies
the same effects as the change-password route (update_password +
revoke_user_refresh_tokens) and drops the stale
app.state.bootstrap_password. MIN_PASSWORD_LENGTH centralised in
auth/storage.py and referenced by the HTTP schema. terminal_prompt.py
carries the pure gate helper (interactive loop stubbed; supplied by the
masked-input module). Also migrates the studio/setup.ps1 launch footer
that still showed the bare wildcard hint.

* README: reconcile remote-access section with opt-in Cloudflare tunnel

* Studio: harden the terminal password gate after review

- run.py: run the gate BEFORE the uvicorn socket binds. On a wildcard
  --cloudflare launch the served HTML injects the bootstrap credential
  for first login, so a pre-gate listener would hand the default
  password to anyone who reaches the raw port while the operator is
  still typing. The gate now also seeds the admin row itself (it can
  run before lifespan startup).
- Headless launches that nothing would protect now fail closed: the
  bootstrap deadline never arms for api-only serving and
  UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT=0 disables it, so warn-and-proceed
  would have promised a shutdown that never comes. Both the CLI and the
  backend refuse to publish in that case; the ordinary headless path
  still warns and relies on the 1h deadline, and no longer auto-fills
  the default credential into HTML served on a public URL.
- storage.update_password gains revoke_refresh_tokens to delete the
  user's refresh tokens in the SAME transaction as the password commit;
  the change-password route and the backend gate use it (a separable
  follow-up delete could fail after the commit and leave a stale
  refresh token able to mint access tokens under the rotated secret).
- clear_bootstrap_password is best-effort: a locked/undeletable file
  must not surface as a failed password change.
- CLI masked reader: disable ISIG like the backend so Ctrl-Z cannot
  suspend the process with the shared terminal stuck in no-echo mode;
  handle Ctrl-C/Ctrl-Z as characters; treat stream EOF mid-line as an
  abort instead of submitting a partial password. Both readers restore
  terminal attrs from a SIGTERM/SIGHUP handler since a finally block
  cannot run when a default-disposition signal terminates the process.
- Backend reader: decode byte-at-a-time through an incremental UTF-8
  decoder so multi-byte characters split across read boundaries are no
  longer dropped; isatty checks tolerate closed/None streams.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: persist bootstrap suppression through lifespan startup

The pre-bind password gate nulled app.state.bootstrap_password, but the
FastAPI lifespan runs after it and re-reads the bootstrap password into
app.state on both admin paths, so a headless public launch could still
serve the injected credential in HTML. Carry a persistent
suppress_bootstrap_injection flag that the lifespan honors instead.

Also drop the quoted Tuple annotation on _terminal_password_gate that
tripped the import-hoist lint (the typing import looked unused).

* Studio CLI: keep the pre-exec auth DB private (0700 dir, 0600 db)

On a fresh install the pre-exposure password gate creates auth/ and
auth.db through the CLI before the backend ever runs, and
sqlite3.connect leaves the DB 0644 under a 022 umask. Mirror backend
storage.get_connection's chmod so the committed password hash and JWT
secret are never world-readable, even if the launch aborts before the
backend applies its own modes.

* Tighten pre-exposure password gate comments

* Studio: delete seeded bootstrap password before headless public re-exec

The headless warn-and-proceed path returns with the default admin
password still active, then re-execs a child Studio process. An old
studio-venv child (mixed-version install) predates the pre-bind gate and
its injection-suppress flag, so its lifespan reads .bootstrap_password
and injects the seeded credential into the public HTML for up to the
bootstrap deadline. A CLI-flag handshake cannot fix this uniformly: the
studio run path uses ignore_unknown_options and an old in-venv child
runs in-process, so it would never reject the flag.

Delete the seeded .bootstrap_password file in the parent before re-exec
so a fresh child of any version reads None and never serves it. This
covers both re-exec paths and both child versions. must_change_password
stays set, so the login page still forces a change and the bootstrap
shutdown timer still arms; only the plaintext-on-disk copy is removed.
Recovery is via a terminal-attached run or reset-password. Backend gate
and CLI warnings updated to match.

* Studio: commit the seeded admin before headless public re-exec

The headless-warn path deletes the seeded .bootstrap_password so a
re-exec'd child cannot inject it, but _ensure_cli_default_admin's INSERT
was never committed and rolled back on conn.close(). On a fresh
STUDIO_HOME an old studio-venv child then found no admin, regenerated a
fresh bootstrap password + file, and injected THAT into the public page,
defeating the deletion.

Commit the seeded admin right after _ensure_cli_default_admin so any
re-exec'd child sees the existing account and does not regenerate.
Regression tests cover both re-exec paths on a fresh (unseeded) DB.

* Studio: fail closed when the bootstrap password file cannot be removed

On the headless public path, deleting .bootstrap_password is the
protection against an old re-exec'd child injecting the seeded
credential. If unlink fails (locked file, read-only auth dir) the file
is still on disk, so warning and proceeding would still leak it for the
bootstrap-timeout window. Abort with a clear error instead. Regression
test covers the unlink-failure fail-closed path.

* Studio: hold no-echo for the whole password line, not per keystroke

The POSIX masked reader set cbreak/no-echo inside _getch_posix and restored
the terminal to echo-on in a finally after every single keystroke, because
_read_password calls _getch once per character. Between one char returning and
the next call re-entering cbreak, ECHO was on, so a keystroke arriving in that
window echoed the password in cleartext.

Move the terminal mode into a _prompt_raw_mode context that _read_password
holds around the entire line (mirroring unsloth_cli/commands/_password_prompt.py,
which already did this), restoring once when the line completes or aborts.
_getch_posix now only reads, since the mode is held by the caller. The context
is a no-op when stdin is not a real terminal, keeping the _getch test seam.

Add a regression test asserting the raw-mode context wraps the read exactly
once and every keystroke is read while it is active.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: strip the seeded bootstrap password when the auth DB check fails

The pre-exposure gate returned early on two auth-DB inspection failures and
proceeded to re-exec without removing the seeded .bootstrap_password:

- _connect_auth_db() failure: a seeded credential from a prior run may still
  be on disk.
- the must_change_password read-back failure: worse, _ensure_cli_default_admin
  had already seeded the admin and the code committed it (writing
  .bootstrap_password) right before the failing SELECT.

In the mixed-version case (a new outer CLI re-execing an old studio-venv child
that predates the pre-bind gate), that child would read the file back and
inject the default admin credential into the public Cloudflare page. The
sibling headless branch already deletes the file for exactly this reason, so
these returns were an inconsistent gap.

Factor the delete-or-fail-closed logic into
_strip_seeded_bootstrap_password_or_exit and call it on both inspection
failures (and reuse it in the headless branch): strip the seeded file first
(version-independent protection), failing closed if the removal itself fails.
must_change_password stays set, so the login page still forces a change and the
bootstrap shutdown timer still arms.

Add tests for both new paths (connect failure and post-commit read-back
failure strip the file and proceed; a failed strip fails closed).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: fail closed when the seeded admin cannot be committed before exposure

The pre-exposure gate wrapped _ensure_cli_default_admin (the INSERT), its
conn.commit(), and the must_change_password read-back in one try, and the
except recovered by stripping .bootstrap_password and proceeding to re-exec on
the assumption the admin was already committed. That assumption only holds when
the failing statement is the SELECT. When the INSERT or the commit itself fails
(e.g. a write lock held past the busy timeout on a fresh install), no admin row
is committed: it rolls back on conn.close(), and a re-exec'd old studio-venv
child (no pre-bind gate) then finds no admin, regenerates a fresh bootstrap
password + file, and serves that default credential on the public Cloudflare
page. Stripping the file cannot stop a regeneration.

Split the seed+commit into its own try that fails closed (refuse the public
launch, best-effort removing any half-written seed file) since we cannot prove a
committed admin; keep the separate read-back failure on the strip-and-proceed
path, where the admin is committed so an old child finds it and will not
regenerate. Add a test for the seed-commit-failure path.

* Studio: decode the CLI masked password reader with errors="replace"

The CLI reader read keystrokes with text-mode sys.stdin.read(1), which raises
UnicodeDecodeError on a pasted non-UTF-8 password (e.g. Latin-1 bytes), or under
PYTHONUTF8 yields a lone surrogate that later crashes the pbkdf2 encode -- either
aborts the launch with a traceback. The backend mirror (terminal_prompt.py)
already reads raw bytes through an incremental decoder with errors="replace".
Mirror that here: read with os.read and an incremental decoder so invalid bytes
map to U+FFFD, iterating over each emitted char (one byte can complete a
replacement plus the next char).

* Studio: resolve the child launcher before the pre-exposure gate

The gate strips the seeded .bootstrap_password on a headless public launch, and
it ran before the re-exec launchability check (studio venv / run.py / console
script present). So a headless launch with an incomplete studio setup would seed
the admin, delete the bootstrap password, then abort because the child could not
be found, leaving the admin at must_change_password=1 with no password ever
shown or injectable: locked out until `unsloth studio reset-password`.

Resolve and validate the child launcher first, in both `studio` (studio_default)
and `studio run`, and only then run the gate, so an unlaunchable setup exits
before anything is stripped. Add a regression test that a missing venv exits
without removing the seeded file.

* Studio: fail closed when the auth DB cannot be opened before exposure

The connect-failure branch of the pre-exposure gate stripped .bootstrap_password
and proceeded, on the assumption a committed admin from a prior run made an old
child find it and not regenerate. But on a fresh public launch whose
_connect_auth_db() itself fails (transient lock during the schema/seed step, or
an unwritable home), no admin is committed, so a mixed-version re-exec child that
predates the backend gate can find no user, generate a fresh bootstrap password,
and serve it on the public Cloudflare page. Stripping a file we cannot vouch for
cannot stop a regeneration.

Make this branch fail closed like the seed/commit failure path: we only continue
past the DB inspection once a committed admin is confirmed. The existing file is
left untouched so a retry (after a transient lock clears) can still prompt.

Update the connect-failure test to assert fail-closed, and give the in-venv
--secure flag test a real STUDIO_HOME with an already-changed admin so the gate
is a no-op rather than relying on a DB-open failure.

* Studio: invalidate seeded bootstrap files before deleting auth.db on reset

reset-password deleted auth.db first, then best-effort unlinked the seeded
.bootstrap_password and desktop secret. unlink() only ignores
FileNotFoundError, so a locked or read-only file (Windows AV, read-only auth
dir) survived while auth.db was gone. The next server start then re-seeded
from that stale plaintext and re-validated the exact credential the reset was
meant to revoke.

Invalidate the credential files first, truncating any that cannot be
unlinked, then delete the DB, so a surviving file can never carry a reusable
secret. clear_bootstrap_password now truncates on unlink failure for the same
reason, and its warning says the contents were cleared rather than claiming
the stale password is already invalid.

* Studio: require a servable frontend before the pre-exposure gate can strip the seeded password

A headless public launch strips the seeded .bootstrap_password before the
re-exec'd child starts. If the child then cannot serve the login page (the only
in-band way to change the seeded password) the admin is locked out
(must_change_password=1, no file, no UI) until reset-password.

Add _require_servable_frontend_or_exit and call it before the gate on both
`unsloth studio` and `unsloth studio run` public launches: fail closed if a
non-api-only public launch has no built frontend dist, before anything is
stripped. A user-supplied --frontend is validated to contain index.html so a
bad path cannot silently bypass the check; an auto-resolved dist is trusted
(_find_frontend_dist already requires index.html) and forwarded to the child.

Model-load aborts on `studio run` remain a residual: the parent must strip for
mixed-version safety (an old studio-venv child has no pre-bind gate) and model
loadability cannot be proven before exec, so that path stays recoverable via
reset-password.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: harden reset-password ordering and validate the in-venv backend before the strip

Three follow-ups to the pre-exposure hardening:

reset-password now deletes auth.db FIRST and proves it is gone before touching
the seeded credential files. If the DB cannot be removed (a running Studio or
Windows holds it open, or a read-only auth dir) it aborts with the credential
files untouched, so a forgotten-password reset is not left half-done with the
recovery credentials deleted while an un-resettable must_change_password=1 DB
survives. After the DB is gone it invalidates the stale credential files
(unlink, else truncate) and fails closed if a file can be neither removed nor
truncated, since a surviving plaintext would be re-seeded by
generate_bootstrap_password() and re-validate the revoked password.

The in-venv (in-process) launch path had no analogue of the re-exec launcher
check: a headless public launch would seed the admin and strip the seeded
.bootstrap_password in the gate before _load_run_module() later failed on a
broken/partial venv, leaving must_change_password=1 with no password to log in.
Add _validate_inproc_backend_before_strip, called on the in-venv path (both
`unsloth studio` and `unsloth studio run`) before the gate on the headless
public path, so a broken backend fails cleanly before anything is stripped. It
is scoped to the headless path so an interactive prompt is not delayed behind a
full backend import.

* Studio: validate the frontend and tunnel before the strip on every public path

Five follow-ups closing the remaining pre-exposure-strip lockouts:

The in-venv (in-process) paths of both `unsloth studio` and `unsloth studio
run` validated the backend but not the frontend before the gate, so a headless
public launch with a missing/bad dist would strip the seeded .bootstrap_password
and then abort in run_server() during frontend setup, leaving
must_change_password=1 with no login page. Both now validate a servable frontend
before the strip (cheap check first, backend import after) and serve the
resolved dist in-process.

The `studio run` re-exec discarded the dist that satisfied the pre-strip check
and only forwarded a user-supplied --frontend. In a shadowed install where the
parent finds a built dist the child cannot, it stripped and exec'd without the
path, and the child aborted during frontend setup. It now forwards the resolved
dist, matching `unsloth studio`.

On a headless --secure launch the bind is loopback, so the Cloudflare tunnel is
the only public exposure. If cloudflared is provably unavailable (found nowhere
and undownloadable) the tunnel cannot start, so stripping the recovery
credential would just lock the user out with no public URL ever served. Add
_tunnel_binary_confirmed_unavailable and, on --secure only, refuse the launch
with the credential preserved rather than strip. Wildcard --cloudflare binds
0.0.0.0 publicly regardless of the tunnel, so it still strips; any uncertainty
(helper not loadable) also still strips, since a possible credential leak
outweighs a recoverable lockout.

clear_bootstrap_password no longer claims it cleared the file's contents when
both unlink and truncate failed; it now reports the stale password is still on
disk and asks the user to remove it manually.

* Studio: fix cloudflared probe path and skip the bootstrap strip for a self-suppressing child

Two follow-ups to the --secure pre-exposure hardening:

The cloudflared availability probe loaded cloudflare_tunnel by file path but not
its backend deps: ensure_cloudflared() -> _cache_path() lazily imports
utils.paths.storage_roots, which only resolves when studio/backend is on
sys.path. From the outer CLI it is not, so the probe saw ensure_cloudflared()
return None (cache unresolvable) and wrongly treated the tunnel as unavailable,
refusing --secure even when cloudflared was cached or downloadable. Add the
backend dir to sys.path for the probe (and remove it after) so the cache path
resolves as it will in the child.

A headless --secure launch stripped the seeded .bootstrap_password before the
child proved the tunnel could actually connect, so a cloudflared that is present
but cannot establish the tunnel (blocked connectivity, Cloudflare outage) left
must_change_password=1 with no recovery credential. But the strip is only needed
when the re-exec'd child is an OLD studio-venv backend with no pre-bind
suppression: this install's own run.py sets app.state.suppress_bootstrap_injection
before binding and never serves the seeded credential publicly. Add
_child_self_suppresses (true in-process, or when the re-exec target is this
install's own run.py by path identity) and skip the strip in that case, keeping
.bootstrap_password as a local recovery credential; the strip stays fully in
force for the studio-venv console-script path and any venv-fallback run.py, where
an old child is actually possible.

* Studio: reword the pre-exposure terminal password prompt

* Studio: warn when -H is overridden by --secure; align pre-exposure prompt wording

- --secure/--secure run: emit a Note (not an error) when -H is a non-loopback
  host, since --secure forces the loopback bind and would otherwise discard -H
  silently.
- Reword the pre-exposure terminal prompt to 'exposed on the public internet'
  in both the backend gate and the CLI mirror.
- Align the CLI success line with the backend ("Password updated for '<user>'.").
- Tests for the new -H warning (present when overridden, absent on loopback).

* Studio: add non-interactive --password to set the initial admin password

Headless hosts (CI, containers, systemd units) have no TTY, so the forced
first-exposure password change could not be completed unattended. Add a
non-interactive way to set the INITIAL admin password before the server binds:

- --password <value>, the UNSLOTH_STUDIO_PASSWORD env var, or --password -
  (read one line from stdin). Off by default; unset falls back to the normal
  interactive terminal prompt / browser setup.
- Applies on any launch (public --secure/--cloudflare or a headless -H 0.0.0.0
  bind), only when the account still has its seeded bootstrap password. An
  already-set password is a hard error, never an override; an invalid value
  (too short, or equal to the bootstrap) fails closed before bind.
- The CLI applies the change in the parent, never forwards --password to the
  re-exec child, and strips UNSLOTH_STUDIO_PASSWORD from the child env so the
  secret never crosses to the child. run.py does the same on the direct path and
  strips the env var so spawned subprocesses (cloudflared, llama-server, tools)
  cannot inherit it.

Mirrors resolve_supplied_password across the CLI and backend, documents the
option in the README (including the argv-visibility caveat), and covers all
flows (env/stdin/literal, fail-closed cases, no-forward, env-strip,
reset-password roundtrip) in the CLI, backend, and unit suites.

* Studio: truncate the stale bootstrap file when unlink fails on a CLI password change

The post-change cleanup in _cli_update_password only warned when
.bootstrap_password could not be unlinked but was still writable (locked file,
read-only auth dir), leaving the old plaintext on disk. If auth.db is later
reset or removed, generate_bootstrap_password() reads that file back and
re-validates the revoked bootstrap password. Truncate the file on unlink
failure so its stale plaintext cannot be re-seeded, mirroring the backend
clear_bootstrap_password(); the password change is already committed, so this
never rolls it back. The warning now states truthfully whether the contents
were cleared or the file must be removed manually.

* Studio: tighten comments

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-15 06:13:25 -07:00

938 lines
31 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""SQLite storage for auth data (user credentials + JWT secret)."""
import hashlib
import hmac
import ipaddress
import os
import secrets
import sqlite3
import threading
from datetime import datetime, timezone
from typing import Optional, Tuple
from utils.paths import auth_db_path, ensure_dir
DB_PATH = auth_db_path()
DEFAULT_ADMIN_USERNAME = "unsloth"
# Single source for the password policy; models/auth.py ChangePasswordRequest
# and the terminal prompt both enforce it. Keep the unsloth_cli mirror in sync.
MIN_PASSWORD_LENGTH = 8
# Plaintext bootstrap password file beside auth.db, deleted on first password
# change so the credential never lingers on disk.
_BOOTSTRAP_PW_PATH = DB_PATH.parent / ".bootstrap_password"
# In-process cache to avoid re-reading the file on every HTML serve.
_bootstrap_password: Optional[str] = None
def generate_bootstrap_password() -> str:
"""Generate a 4-word diceware passphrase and persist it to disk.
Persisted (the DB stores only the hash) so it survives restarts; later
calls return the persisted value.
"""
global _bootstrap_password
# Cached in this process?
if _bootstrap_password is not None:
return _bootstrap_password
# Persisted from a previous run?
if _BOOTSTRAP_PW_PATH.is_file():
_bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
if _bootstrap_password:
return _bootstrap_password
# First startup: generate a fresh passphrase.
import diceware
_bootstrap_password = diceware.get_passphrase(
options = diceware.handle_options(args = ["-n", "4", "-d", "", "-c"])
)
# Persist so the same passphrase survives restarts until password change.
ensure_dir(_BOOTSTRAP_PW_PATH.parent)
_BOOTSTRAP_PW_PATH.write_text(_bootstrap_password)
try:
os.chmod(_BOOTSTRAP_PW_PATH, 0o600)
except OSError:
pass
return _bootstrap_password
def get_bootstrap_password() -> Optional[str]:
"""Return the cached bootstrap password, or None if not yet generated."""
return _bootstrap_password
def _load_bootstrap_password() -> Optional[str]:
"""Load an existing bootstrap password without creating one."""
global _bootstrap_password
_bootstrap_password = None
if _BOOTSTRAP_PW_PATH.is_file():
bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
if bootstrap_password:
_bootstrap_password = bootstrap_password
return _bootstrap_password
def clear_bootstrap_password() -> None:
"""Delete the persisted bootstrap password file (after a password change).
Best-effort: the new hash is already committed, so a locked/undeletable file
(Windows AV, read-only auth dir) must not fail the change.
"""
global _bootstrap_password
_bootstrap_password = None
if _BOOTSTRAP_PW_PATH.is_file():
try:
_BOOTSTRAP_PW_PATH.unlink(missing_ok = True)
except OSError as e:
# Removal failed (Windows AV, read-only auth dir). The hash is already
# committed, so don't fail the change -- but truncate the file so its
# stale plaintext can't be re-seeded by generate_bootstrap_password()
# if a later reset-password deletes auth.db and re-validates it.
try:
_BOOTSTRAP_PW_PATH.write_text("")
cleared = True
except OSError:
cleared = False
import sys
if cleared:
message = (
f"Warning: could not delete {_BOOTSTRAP_PW_PATH.name} ({e}); "
"cleared its contents so the old bootstrap password cannot be reused."
)
else:
# Neither removed nor truncated: stale plaintext is still on disk
# and would be reused if auth.db is reset. Don't claim otherwise.
message = (
f"Warning: could not delete or clear {_BOOTSTRAP_PW_PATH.name} ({e}); "
"its old bootstrap password is still on disk. Remove it manually to "
"prevent reuse after a reset."
)
print(message, file = sys.stderr, flush = True)
def _hash_token(token: str) -> str:
"""SHA-256 hash helper for refresh token storage.
Plain SHA-256 is intentional: refresh tokens are 384-bit random strings, so
a slow KDF adds no security while costing per-refresh latency. API keys use
the separate ``_pbkdf2_api_key`` helper, only to satisfy CodeQL's
``py/weak-sensitive-data-hashing`` query, not for crypto reasons.
"""
return hashlib.sha256(token.encode("utf-8")).hexdigest()
def get_connection() -> sqlite3.Connection:
"""Get a connection to the auth database, creating tables if needed."""
ensure_dir(DB_PATH.parent)
conn = sqlite3.connect(DB_PATH)
# Keep the auth dir + DB private (they hold the JWT/identity secrets and
# password hashes); sqlite3.connect would otherwise create the DB 0644 under
# a 022 umask, letting another OS user read the identity secret and forge proofs.
for _path, _mode in ((DB_PATH.parent, 0o700), (DB_PATH, 0o600)):
try:
os.chmod(_path, _mode)
except OSError:
pass
conn.row_factory = sqlite3.Row
# WAL lets token reads run concurrently with refresh-token writes;
# busy_timeout bounds lock waits. Matches the other Studio SQLite stores.
# Set busy_timeout first: switching journal_mode needs a lock, so if a
# refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY;
# with busy_timeout already in effect it waits instead of failing and leaving
# this connection on SQLite's default zero lock wait.
try:
conn.execute("PRAGMA busy_timeout=5000")
conn.execute("PRAGMA journal_mode=WAL")
except sqlite3.Error:
pass
conn.execute(
"""
CREATE TABLE IF NOT EXISTS auth_user (
id INTEGER PRIMARY KEY,
username TEXT UNIQUE NOT NULL,
password_salt TEXT NOT NULL,
password_hash TEXT NOT NULL,
jwt_secret TEXT NOT NULL,
must_change_password INTEGER NOT NULL DEFAULT 0
);
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS refresh_tokens (
id INTEGER PRIMARY KEY,
token_hash TEXT NOT NULL,
username TEXT NOT NULL,
expires_at TEXT NOT NULL,
is_desktop INTEGER NOT NULL DEFAULT 0
);
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS api_keys (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT NOT NULL,
key_prefix TEXT NOT NULL,
key_hash TEXT NOT NULL UNIQUE,
name TEXT NOT NULL DEFAULT '',
created_at TEXT NOT NULL,
last_used_at TEXT,
expires_at TEXT,
is_active INTEGER NOT NULL DEFAULT 1,
is_internal INTEGER NOT NULL DEFAULT 0
);
"""
)
api_key_columns = {row["name"] for row in conn.execute("PRAGMA table_info(api_keys)")}
if "is_internal" not in api_key_columns:
conn.execute("ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS app_secrets (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);
"""
)
columns = {row["name"] for row in conn.execute("PRAGMA table_info(auth_user)")}
if "must_change_password" not in columns:
conn.execute(
"ALTER TABLE auth_user ADD COLUMN must_change_password INTEGER NOT NULL DEFAULT 0"
)
refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")}
if "is_desktop" not in refresh_columns:
conn.execute("ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0")
conn.commit()
return conn
# ── API-key PBKDF2 salt ────────────────────────────────────────────────
#
# Module-level cache for the persistent API-key PBKDF2 salt, populated lazily
# via ``_get_or_create_api_key_pbkdf2_salt``. No lock needed: (a) ``INSERT OR
# IGNORE`` is atomic at the SQLite layer and (b) concurrent populations
# converge on the same value, so the worst case is a harmless duplicate read
# on startup.
_api_key_pbkdf2_salt_cache: Optional[bytes] = None
def _get_or_create_api_key_pbkdf2_salt() -> bytes:
"""Return the persistent API-key PBKDF2 salt, generating it once if missing.
Hex-encoded 32-byte random value in ``app_secrets``. Regenerated only when
the row is missing (fresh install, or operator deleted it).
"""
global _api_key_pbkdf2_salt_cache
if _api_key_pbkdf2_salt_cache is not None:
return _api_key_pbkdf2_salt_cache
conn = get_connection()
try:
cur = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
("api_key_pbkdf2_salt",),
)
row = cur.fetchone()
if row is None:
new_value = secrets.token_hex(32) # 32 bytes -> 64 hex chars
conn.execute(
"INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)",
("api_key_pbkdf2_salt", new_value),
)
conn.commit()
cur = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
("api_key_pbkdf2_salt",),
)
row = cur.fetchone()
salt = bytes.fromhex(row["value"])
finally:
conn.close()
_api_key_pbkdf2_salt_cache = salt
return salt
# Secret answering the /api/auth/identity challenge (HMAC(secret, nonce)). Lives
# in this same-user DB so a port squatter or remote/fake server can't forge a
# proof. Separate from the per-user JWT secret.
_IDENTITY_SECRET_DB_KEY = "studio_identity_secret"
_identity_secret_cache: Optional[bytes] = None
def get_or_create_identity_secret() -> bytes:
"""Return the identity secret (hex 32-byte row in app_secrets), creating it once."""
global _identity_secret_cache
if _identity_secret_cache is not None:
return _identity_secret_cache
conn = get_connection()
try:
row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_IDENTITY_SECRET_DB_KEY,),
).fetchone()
if row is None:
conn.execute(
"INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)",
(_IDENTITY_SECRET_DB_KEY, secrets.token_hex(32)),
)
conn.commit()
row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_IDENTITY_SECRET_DB_KEY,),
).fetchone()
secret = bytes.fromhex(row["value"])
finally:
conn.close()
_identity_secret_cache = secret
return secret
def compute_identity_proof(nonce: bytes, host: str, port: int) -> str:
"""HMAC-SHA256 proof that the caller holds this install's identity secret,
bound to the loopback address and port the connection landed on. A proof
relayed from a Studio on a different address/port (a squatter proxying to the
real one, e.g. localhost resolving to ::1 while Studio is on 127.0.0.1) was
computed for that other endpoint and won't match the one the client dialed."""
try:
host = ipaddress.ip_address(host).compressed # normalise 127.0.0.1 / ::1 forms
except ValueError:
host = (host or "").lower()
msg = b"|".join([nonce, host.encode(), str(int(port)).encode()])
return hmac.new(get_or_create_identity_secret(), msg, hashlib.sha256).hexdigest()
# Capability secret for public ``/p`` preview share links. HMAC(secret, ref)
# turns the deterministic preview ref into an unguessable bearer capability, so a
# guessed run/checkpoint name can't reach inference. Dedicated (not the per-user
# JWT secret) so rotating it revokes every shared link without touching logins.
_PREVIEW_LINK_SECRET_DB_KEY = "preview_link_secret"
_preview_link_secret_cache: Optional[bytes] = None
def get_or_create_preview_link_secret() -> bytes:
"""Return the preview-link signing secret (hex 32-byte row in app_secrets), creating it once."""
global _preview_link_secret_cache
if _preview_link_secret_cache is not None:
return _preview_link_secret_cache
conn = get_connection()
try:
row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_PREVIEW_LINK_SECRET_DB_KEY,),
).fetchone()
if row is None:
conn.execute(
"INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)",
(_PREVIEW_LINK_SECRET_DB_KEY, secrets.token_hex(32)),
)
conn.commit()
row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_PREVIEW_LINK_SECRET_DB_KEY,),
).fetchone()
secret = bytes.fromhex(row["value"])
finally:
conn.close()
_preview_link_secret_cache = secret
return secret
def rotate_preview_link_secret() -> bytes:
"""Rotate the preview-link secret, immediately revoking every outstanding ``/p`` share link."""
global _preview_link_secret_cache
new_secret_hex = secrets.token_hex(32)
conn = get_connection()
try:
conn.execute(
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
(_PREVIEW_LINK_SECRET_DB_KEY, new_secret_hex),
)
conn.commit()
finally:
conn.close()
secret = bytes.fromhex(new_secret_hex)
_preview_link_secret_cache = secret
return secret
_API_KEY_PBKDF2_ITERATIONS = 100_000
DESKTOP_SECRET_PREFIX = "desktop-"
_DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash"
_DESKTOP_SECRET_CREATED_AT_KEY = "desktop_secret_created_at"
def _pbkdf2_api_key(raw_key: str) -> str:
"""PBKDF2-HMAC-SHA256 an API key with a persistent server-side salt.
For API-key storage ONLY, not refresh tokens. The slow KDF is only to
appease CodeQL's ``py/weak-sensitive-data-hashing`` query, not a crypto
requirement (API keys are random 128-bit tokens). The salt lives in
``app_secrets`` so dumping ``api_keys`` alone can't derive hashes.
"""
salt = _get_or_create_api_key_pbkdf2_salt()
dk = hashlib.pbkdf2_hmac(
"sha256",
raw_key.encode("utf-8"),
salt,
_API_KEY_PBKDF2_ITERATIONS,
)
return dk.hex()
def _pbkdf2_desktop_secret(raw_secret: str) -> str:
return _pbkdf2_api_key(raw_secret)
# Memoize the deterministic raw-key -> PBKDF2-hash derivation so the 100k-round
# KDF runs once per key instead of on every authenticated request. Keyed by a
# salted HMAC of the key (not the key itself); revocation/expiry are still
# enforced by the SQLite read on every call, so a cache hit only skips the KDF.
# Only keys present in the DB are cached, so unknown-key spam can't grow it.
_api_key_hash_cache: dict[str, str] = {}
_API_KEY_HASH_CACHE_MAX = 4096
_api_key_hash_cache_lock = threading.Lock()
def _api_key_cache_id(raw_key: str) -> str:
"""Cache id for a raw key: salted HMAC-SHA256 (not the key itself)."""
return hmac.new(
_get_or_create_api_key_pbkdf2_salt(), raw_key.encode("utf-8"), hashlib.sha256
).hexdigest()
def _reset_api_key_hash_cache() -> None:
"""Drop memoized derivations (tests / salt change)."""
with _api_key_hash_cache_lock:
_api_key_hash_cache.clear()
def is_initialized() -> bool:
"""Check if auth is ready for login (at least one user exists in DB)."""
conn = get_connection()
cur = conn.execute("SELECT COUNT(*) AS c FROM auth_user")
row = cur.fetchone()
conn.close()
return bool(row["c"])
def create_initial_user(
username: str,
password: str,
jwt_secret: str,
*,
must_change_password: bool = False,
) -> None:
"""
Create the initial admin user in the database.
Raises sqlite3.IntegrityError if username already exists.
"""
from .hashing import hash_password
salt, pwd_hash = hash_password(password)
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO auth_user (
username,
password_salt,
password_hash,
jwt_secret,
must_change_password
)
VALUES (?, ?, ?, ?, ?)
""",
(username, salt, pwd_hash, jwt_secret, int(must_change_password)),
)
conn.commit()
finally:
conn.close()
def delete_user(username: str) -> None:
"""
Delete a user from the database.
Used for rollback when user creation fails partway through bootstrap.
"""
conn = get_connection()
try:
conn.execute("DELETE FROM auth_user WHERE username = ?", (username,))
conn.commit()
finally:
conn.close()
def get_user_and_secret(username: str) -> Optional[Tuple[str, str, str, bool]]:
"""
Get user's password salt, hash, and JWT secret.
Returns (password_salt, password_hash, jwt_secret, must_change_password)
or None if user not found.
"""
conn = get_connection()
try:
cur = conn.execute(
"""
SELECT password_salt, password_hash, jwt_secret, must_change_password
FROM auth_user
WHERE username = ?
""",
(username,),
)
row = cur.fetchone()
if not row:
return None
return (
row["password_salt"],
row["password_hash"],
row["jwt_secret"],
bool(row["must_change_password"]),
)
finally:
conn.close()
def get_jwt_secret(username: str) -> Optional[str]:
"""Return the current JWT signing secret for a user."""
conn = get_connection()
try:
cur = conn.execute(
"SELECT jwt_secret FROM auth_user WHERE username = ?",
(username,),
)
row = cur.fetchone()
return row["jwt_secret"] if row else None
finally:
conn.close()
def requires_password_change(username: str) -> bool:
"""Return whether the user must change the seeded default password."""
conn = get_connection()
try:
cur = conn.execute(
"SELECT must_change_password FROM auth_user WHERE username = ?",
(username,),
)
row = cur.fetchone()
return bool(row and row["must_change_password"])
finally:
conn.close()
def load_jwt_secret() -> str:
"""
Load the JWT secret from the database.
Raises RuntimeError if no auth user has been created yet.
"""
conn = get_connection()
try:
cur = conn.execute("SELECT jwt_secret FROM auth_user LIMIT 1")
row = cur.fetchone()
if not row:
raise RuntimeError(
"Auth is not initialized. Wait for the seeded admin bootstrap to complete."
)
return row["jwt_secret"]
finally:
conn.close()
def ensure_default_admin() -> bool:
"""Seed the default admin account on first startup.
Uses a randomly generated diceware passphrase as the bootstrap password.
Returns True when the default admin was created in this call.
"""
if get_user_and_secret(DEFAULT_ADMIN_USERNAME) is not None:
_load_bootstrap_password()
return False
bootstrap_pw = generate_bootstrap_password()
try:
create_initial_user(
username = DEFAULT_ADMIN_USERNAME,
password = bootstrap_pw,
jwt_secret = secrets.token_urlsafe(64),
must_change_password = True,
)
return True
except sqlite3.IntegrityError:
return False
def update_password(
username: str,
new_password: str,
*,
revoke_refresh_tokens: bool = False,
) -> bool:
"""Update password, clear first-login requirement, rotate JWT secret.
``revoke_refresh_tokens`` deletes the user's refresh tokens in the SAME
transaction: a separate delete could fail after the password commit and
leave a pre-change token still able to mint access tokens.
"""
from .hashing import hash_password
salt, pwd_hash = hash_password(new_password)
jwt_secret = secrets.token_urlsafe(64)
conn = get_connection()
try:
cursor = conn.execute(
"""
UPDATE auth_user
SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0
WHERE username = ?
""",
(salt, pwd_hash, jwt_secret, username),
)
if revoke_refresh_tokens and cursor.rowcount > 0:
conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,))
conn.commit()
if cursor.rowcount > 0:
clear_bootstrap_password()
clear_desktop_secret()
return cursor.rowcount > 0
finally:
conn.close()
def save_refresh_token(
token: str,
username: str,
expires_at: str,
*,
is_desktop: bool = False,
) -> None:
"""
Store a hashed refresh token with its associated username and expiry.
"""
token_hash = _hash_token(token)
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO refresh_tokens (token_hash, username, expires_at, is_desktop)
VALUES (?, ?, ?, ?)
""",
(token_hash, username, expires_at, int(is_desktop)),
)
conn.commit()
finally:
conn.close()
def consume_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
"""Atomically validate-and-delete a refresh token for single-use rotation.
DELETE RETURNING fuses validate and delete into one statement so two
concurrent refresh requests cannot both consume the same token.
"""
token_hash = _hash_token(token)
now = datetime.now(timezone.utc).isoformat()
conn = get_connection()
try:
conn.execute(
"DELETE FROM refresh_tokens WHERE expires_at < ?",
(now,),
)
cur = conn.execute(
"""
DELETE FROM refresh_tokens
WHERE token_hash = ? AND expires_at >= ?
RETURNING username, is_desktop
""",
(token_hash, now),
)
row = cur.fetchone()
conn.commit()
if row is None:
return None
return row["username"], bool(row["is_desktop"])
finally:
conn.close()
def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
"""
Verify a refresh token and return the username plus desktop marker.
Returns the username and desktop marker if valid and not expired, None otherwise.
The token is NOT consumed — it stays valid until it expires.
"""
token_hash = _hash_token(token)
conn = get_connection()
try:
# Opportunistically clean up expired tokens
conn.execute(
"DELETE FROM refresh_tokens WHERE expires_at < ?",
(datetime.now(timezone.utc).isoformat(),),
)
conn.commit()
cur = conn.execute(
"""
SELECT id, username, expires_at, is_desktop FROM refresh_tokens
WHERE token_hash = ?
""",
(token_hash,),
)
row = cur.fetchone()
if row is None:
return None
# Check expiry
expires_at = datetime.fromisoformat(row["expires_at"])
if datetime.now(timezone.utc) > expires_at:
conn.execute("DELETE FROM refresh_tokens WHERE id = ?", (row["id"],))
conn.commit()
return None
return row["username"], bool(row["is_desktop"])
finally:
conn.close()
def revoke_user_refresh_tokens(username: str) -> None:
"""Revoke all refresh tokens for a user (e.g. on logout)."""
conn = get_connection()
try:
conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,))
conn.commit()
finally:
conn.close()
def create_desktop_secret() -> str:
"""Create/rotate the local desktop credential and return it once."""
ensure_default_admin()
raw_secret = DESKTOP_SECRET_PREFIX + secrets.token_urlsafe(48)
secret_hash = _pbkdf2_desktop_secret(raw_secret)
now = datetime.now(timezone.utc).isoformat()
conn = get_connection()
try:
conn.execute(
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
(_DESKTOP_SECRET_HASH_KEY, secret_hash),
)
conn.execute(
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
(_DESKTOP_SECRET_CREATED_AT_KEY, now),
)
conn.commit()
return raw_secret
finally:
conn.close()
def validate_desktop_secret(raw_secret: str) -> Optional[str]:
"""Return the real admin username when the desktop secret matches."""
if not raw_secret.startswith(DESKTOP_SECRET_PREFIX):
return None
if get_user_and_secret(DEFAULT_ADMIN_USERNAME) is None:
return None
secret_hash = _pbkdf2_desktop_secret(raw_secret)
conn = get_connection()
try:
cur = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_DESKTOP_SECRET_HASH_KEY,),
)
row = cur.fetchone()
if row is None:
return None
if not secrets.compare_digest(row["value"], secret_hash):
return None
return DEFAULT_ADMIN_USERNAME
finally:
conn.close()
def clear_desktop_secret() -> None:
"""Remove backend-side desktop auth state."""
conn = get_connection()
try:
conn.execute(
"DELETE FROM app_secrets WHERE key IN (?, ?)",
(_DESKTOP_SECRET_HASH_KEY, _DESKTOP_SECRET_CREATED_AT_KEY),
)
conn.commit()
finally:
conn.close()
# ---------------------------------------------------------------------------
# API key management
# ---------------------------------------------------------------------------
API_KEY_PREFIX = "sk-unsloth-"
def create_api_key(
username: str,
name: str,
expires_at: Optional[str] = None,
internal: bool = False,
) -> Tuple[str, dict]:
"""Create a new API key for *username*.
Returns ``(raw_key, row_dict)`` where *raw_key* is shown to the user
exactly once. The database only stores the PBKDF2 hash.
Pass ``internal=True`` for keys minted by workflows (e.g. data-recipe
runs) that should not appear in user-facing key listings.
"""
raw_key = API_KEY_PREFIX + secrets.token_hex(16)
key_hash = _pbkdf2_api_key(raw_key)
key_prefix = raw_key[len(API_KEY_PREFIX) : len(API_KEY_PREFIX) + 8]
now = datetime.now(timezone.utc).isoformat()
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at, is_internal)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(
username,
key_prefix,
key_hash,
name,
now,
expires_at,
1 if internal else 0,
),
)
conn.commit()
cur = conn.execute("SELECT * FROM api_keys WHERE key_hash = ?", (key_hash,))
row = cur.fetchone()
return raw_key, dict(row)
finally:
conn.close()
def list_api_keys(username: str, include_internal: bool = False) -> list:
"""Return API keys for *username*. Internal workflow keys are hidden
by default so they do not clutter user-facing UIs."""
conn = get_connection()
try:
if include_internal:
cur = conn.execute(
"""
SELECT id, username, key_prefix, name, created_at, last_used_at,
expires_at, is_active, is_internal
FROM api_keys
WHERE username = ?
ORDER BY created_at DESC
""",
(username,),
)
else:
cur = conn.execute(
"""
SELECT id, username, key_prefix, name, created_at, last_used_at,
expires_at, is_active, is_internal
FROM api_keys
WHERE username = ? AND is_internal = 0
ORDER BY created_at DESC
""",
(username,),
)
return [dict(row) for row in cur.fetchall()]
finally:
conn.close()
def revoke_api_key(username: str, key_id: int) -> bool:
"""Soft-delete an API key. Returns True if a matching row was found."""
conn = get_connection()
try:
cursor = conn.execute(
"UPDATE api_keys SET is_active = 0 WHERE id = ? AND username = ?",
(key_id, username),
)
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def revoke_internal_api_key(key_id: int) -> bool:
"""Revoke an internal workflow-minted key without requiring a username.
Used by the recipe runner to retire its sk-unsloth-* key once the job
terminates, shrinking the window a leaked key could be abused.
"""
conn = get_connection()
try:
cursor = conn.execute(
"UPDATE api_keys SET is_active = 0 WHERE id = ? AND is_internal = 1",
(key_id,),
)
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def validate_api_key(raw_key: str) -> Optional[str]:
"""Validate *raw_key* and return the owning username, or ``None``.
Also updates ``last_used_at`` on success.
"""
cache_id = _api_key_cache_id(raw_key)
cached_hash = _api_key_hash_cache.get(cache_id)
key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key)
conn = get_connection()
try:
cur = conn.execute(
"SELECT id, username, is_active, expires_at FROM api_keys WHERE key_hash = ?",
(key_hash,),
)
row = cur.fetchone()
if row is None:
return None
# Real key: memoize so later requests skip the KDF. Bounded; clear on overflow.
if cached_hash is None:
with _api_key_hash_cache_lock:
if len(_api_key_hash_cache) >= _API_KEY_HASH_CACHE_MAX:
_api_key_hash_cache.clear()
_api_key_hash_cache[cache_id] = key_hash
if not row["is_active"]:
return None
if row["expires_at"] is not None:
expires = datetime.fromisoformat(row["expires_at"])
if datetime.now(timezone.utc) > expires:
return None
conn.execute(
"UPDATE api_keys SET last_used_at = ? WHERE id = ?",
(datetime.now(timezone.utc).isoformat(), row["id"]),
)
conn.commit()
return row["username"]
finally:
conn.close()