243 lines
12 KiB
Python
243 lines
12 KiB
Python
"""SQLite persistence. No credential-bearing row is ever a public API schema."""
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import hmac
|
||
import os
|
||
import re
|
||
import secrets
|
||
import sqlite3
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
PUBLIC_FIELDS = ('id', 'username', 'display_name', 'role', 'status', 'container_name', 'error')
|
||
SLUG = re.compile(r'[a-z][a-z0-9]*(?:-[a-z0-9]+)*\Z')
|
||
MODEL = re.compile(r'[!-~]{1,200}\Z')
|
||
THINKING_LEVELS = ('none', 'minimal', 'low', 'medium', 'high', 'xhigh')
|
||
|
||
|
||
def public_user(user):
|
||
return {key: user[key] for key in PUBLIC_FIELDS}
|
||
|
||
|
||
def validate_username(value):
|
||
if not isinstance(value, str) or not 1 <= len(value) <= 48 or not SLUG.fullmatch(value):
|
||
raise ValueError('Username must be a lowercase slug (1–48 characters)')
|
||
return value
|
||
|
||
|
||
def validate_password(value):
|
||
if not isinstance(value, str) or not 12 <= len(value) <= 256 or len(value.encode('utf-8')) > 1024:
|
||
raise ValueError('Password must contain 12–256 characters')
|
||
return value
|
||
|
||
|
||
def password_hash(password):
|
||
salt = secrets.token_bytes(16)
|
||
digest = hashlib.scrypt(password.encode(), salt=salt, n=16384, r=8, p=1, dklen=32)
|
||
return f'scrypt$16384$8$1${salt.hex()}${digest.hex()}'
|
||
|
||
|
||
def password_matches(password, encoded):
|
||
try:
|
||
_, n, r, p, salt, digest = encoded.split('$')
|
||
actual = hashlib.scrypt(password.encode(), salt=bytes.fromhex(salt), n=int(n), r=int(r), p=int(p), dklen=32)
|
||
return hmac.compare_digest(actual.hex(), digest)
|
||
except (ValueError, TypeError):
|
||
return False
|
||
|
||
|
||
def token_hash(token):
|
||
return hashlib.sha256(token.encode()).hexdigest()
|
||
|
||
|
||
class Accounts:
|
||
def __init__(self, path):
|
||
if path != ':memory:':
|
||
Path(path).parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||
self.lock = threading.RLock()
|
||
if path != ':memory:':
|
||
# Create privately before SQLite can create journal companions.
|
||
fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o600)
|
||
os.close(fd)
|
||
os.chmod(path, 0o600)
|
||
self.db = sqlite3.connect(path, check_same_thread=False)
|
||
self.db.row_factory = sqlite3.Row
|
||
self.db.execute('PRAGMA foreign_keys=ON')
|
||
self.db.execute('PRAGMA busy_timeout=5000')
|
||
self.db.execute('PRAGMA journal_mode=WAL')
|
||
self.db.executescript('''
|
||
CREATE TABLE IF NOT EXISTS users (
|
||
id TEXT PRIMARY KEY, username TEXT NOT NULL UNIQUE,
|
||
display_name TEXT NOT NULL, role TEXT NOT NULL,
|
||
status TEXT NOT NULL, container_name TEXT NOT NULL,
|
||
error TEXT, password_hash TEXT NOT NULL,
|
||
hermes_url TEXT, hermes_api_key TEXT, desired TEXT NOT NULL DEFAULT 'ready'
|
||
);
|
||
CREATE TABLE IF NOT EXISTS account_model_settings (
|
||
user_id TEXT PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE,
|
||
model TEXT NOT NULL, thinking_level TEXT NOT NULL
|
||
);
|
||
CREATE TABLE IF NOT EXISTS sessions (
|
||
hash TEXT PRIMARY KEY, user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||
expires_at REAL NOT NULL
|
||
);
|
||
CREATE TABLE IF NOT EXISTS rooms (
|
||
room_name TEXT PRIMARY KEY, user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||
expires_at REAL NOT NULL, session_hash TEXT
|
||
);
|
||
''')
|
||
if 'session_hash' not in {r[1] for r in self.db.execute('PRAGMA table_info(rooms)')}:
|
||
self.db.execute('ALTER TABLE rooms ADD COLUMN session_hash TEXT')
|
||
self.db.commit()
|
||
if path != ':memory:':
|
||
os.chmod(path, 0o600)
|
||
for suffix in ('-wal', '-shm'):
|
||
companion = Path(path + suffix)
|
||
if companion.exists():
|
||
os.chmod(companion, 0o600)
|
||
|
||
def close(self):
|
||
self.db.close()
|
||
|
||
def users(self):
|
||
with self.lock:
|
||
return [dict(r) for r in self.db.execute('SELECT * FROM users ORDER BY username')]
|
||
|
||
def get(self, user_id):
|
||
with self.lock:
|
||
row = self.db.execute('SELECT * FROM users WHERE id=?', (user_id,)).fetchone()
|
||
return dict(row) if row else None
|
||
|
||
def get_model_settings(self, user_id):
|
||
with self.lock:
|
||
# The join checks account availability and reads settings in one
|
||
# snapshot, including accounts that have not finished provisioning.
|
||
row = self.db.execute('''
|
||
SELECT COALESCE(s.model, 'auto') AS model,
|
||
COALESCE(s.thinking_level, 'none') AS thinking_level
|
||
FROM users u LEFT JOIN account_model_settings s ON s.user_id=u.id
|
||
WHERE u.id=? AND u.desired!='deleted'
|
||
''', (user_id,)).fetchone()
|
||
if row is None:
|
||
raise KeyError(user_id)
|
||
return dict(row)
|
||
|
||
def set_model_settings(self, user_id, model, thinking_level):
|
||
if not isinstance(model, str) or not MODEL.fullmatch(model):
|
||
raise ValueError('Model must contain 1–200 printable ASCII characters without whitespace')
|
||
if not isinstance(thinking_level, str) or thinking_level not in THINKING_LEVELS:
|
||
raise ValueError('Invalid thinking level')
|
||
with self.lock, self.db:
|
||
# Guard and upsert are one write statement: a delete on another
|
||
# connection cannot interleave between checking and storing.
|
||
cursor = self.db.execute('''
|
||
INSERT INTO account_model_settings(user_id,model,thinking_level)
|
||
SELECT id,?,? FROM users WHERE id=? AND desired!='deleted'
|
||
ON CONFLICT(user_id) DO UPDATE SET
|
||
model=excluded.model, thinking_level=excluded.thinking_level
|
||
''', (model, thinking_level, user_id))
|
||
if cursor.rowcount != 1:
|
||
raise KeyError(user_id)
|
||
return {'model': model, 'thinking_level': thinking_level}
|
||
|
||
def by_username(self, username):
|
||
with self.lock:
|
||
row = self.db.execute('SELECT * FROM users WHERE username=?', (username,)).fetchone()
|
||
return dict(row) if row else None
|
||
|
||
def create(self, username, display_name, role, password):
|
||
validate_username(username)
|
||
validate_password(password)
|
||
if role not in ('admin', 'user'):
|
||
raise ValueError('Invalid role')
|
||
if not isinstance(display_name, str) or not 1 <= len(display_name.strip()) <= 100:
|
||
raise ValueError('Display name must contain 1–100 characters')
|
||
user_id = uuid.uuid4().hex
|
||
hashed = password_hash(password)
|
||
with self.lock, self.db:
|
||
self.db.execute('INSERT INTO users(id,username,display_name,role,status,container_name,password_hash) VALUES(?,?,?,?,?,?,?)',
|
||
(user_id, username, display_name.strip(), role, 'provisioning', f'hope-hermes-{username}-{user_id[:8]}', hashed))
|
||
return self.get(user_id)
|
||
|
||
def update(self, user_id, **values):
|
||
allowed = {'status', 'error', 'hermes_url', 'hermes_api_key', 'container_name'}
|
||
if not values or not set(values) <= allowed:
|
||
raise ValueError('Invalid account update')
|
||
with self.lock, self.db:
|
||
guard = " AND desired!='deleted'" if values.get('status') in ('ready', 'provisioning', 'error') else ''
|
||
self.db.execute('UPDATE users SET ' + ','.join(f'{k}=?' for k in values) + ' WHERE id=?' + guard, (*values.values(), user_id))
|
||
|
||
def begin_delete(self, user_id):
|
||
with self.lock, self.db:
|
||
user = self.get(user_id)
|
||
if user is None:
|
||
raise KeyError(user_id)
|
||
if user['role'] == 'admin' and user['desired'] != 'deleted':
|
||
count = self.db.execute("SELECT count(*) FROM users WHERE role='admin' AND desired!='deleted'").fetchone()[0]
|
||
if count <= 1:
|
||
raise ValueError('Cannot delete the last administrator')
|
||
self.db.execute("UPDATE users SET desired='deleted',status='deleting',error=NULL WHERE id=?", (user_id,))
|
||
self.db.execute('DELETE FROM sessions WHERE user_id=?', (user_id,))
|
||
self.db.execute('DELETE FROM rooms WHERE user_id=?', (user_id,))
|
||
return self.get(user_id)
|
||
|
||
def finish_delete(self, user_id):
|
||
with self.lock, self.db:
|
||
self.db.execute("DELETE FROM users WHERE id=? AND desired='deleted'", (user_id,))
|
||
|
||
def reset_password(self, user_id, password):
|
||
hashed = password_hash(validate_password(password))
|
||
with self.lock, self.db:
|
||
user = self.get(user_id)
|
||
if user is None or user['desired'] == 'deleted':
|
||
raise KeyError(user_id)
|
||
self.db.execute('UPDATE users SET password_hash=? WHERE id=?', (hashed, user_id))
|
||
self.db.execute('DELETE FROM sessions WHERE user_id=?', (user_id,))
|
||
self.db.execute('DELETE FROM rooms WHERE user_id=?', (user_id,))
|
||
|
||
def issue_session(self, user_id, ttl, expected_password_hash):
|
||
token = secrets.token_urlsafe(32)
|
||
expires = time.time() + ttl
|
||
with self.lock, self.db:
|
||
self.db.execute('DELETE FROM sessions WHERE expires_at<=?', (time.time(),))
|
||
# Single SQL statement atomically compares the verified credential
|
||
# and inserts; reset/delete cannot interleave with this check.
|
||
cursor = self.db.execute("INSERT INTO sessions(hash,user_id,expires_at) SELECT ?,id,? FROM users WHERE id=? AND desired!='deleted' AND password_hash=?",
|
||
(token_hash(token), expires, user_id, expected_password_hash))
|
||
if cursor.rowcount != 1:
|
||
raise ValueError('Account credentials changed or unavailable')
|
||
return token, expires
|
||
|
||
def resolve_session(self, token):
|
||
if not isinstance(token, str) or not 20 <= len(token) <= 256:
|
||
return None
|
||
with self.lock:
|
||
row = self.db.execute("SELECT u.* FROM sessions s JOIN users u ON u.id=s.user_id WHERE s.hash=? AND s.expires_at>? AND u.desired!='deleted'", (token_hash(token), time.time())).fetchone()
|
||
return dict(row) if row else None
|
||
|
||
def revoke(self, token):
|
||
with self.lock, self.db:
|
||
self.db.execute('DELETE FROM sessions WHERE hash=?', (token_hash(token),))
|
||
|
||
def grant_room(self, user_id, room_name, ttl, session_token):
|
||
with self.lock, self.db:
|
||
user = self.get(user_id)
|
||
if not user or user['desired'] == 'deleted' or user['status'] != 'ready':
|
||
raise ValueError('Account not ready')
|
||
if not isinstance(session_token, str) or not 20 <= len(session_token) <= 256:
|
||
raise ValueError('Invalid session')
|
||
hashed = token_hash(session_token)
|
||
session = self.db.execute('SELECT expires_at FROM sessions WHERE hash=? AND user_id=? AND expires_at>?', (hashed, user_id, time.time())).fetchone()
|
||
if session is None:
|
||
raise ValueError('Invalid session')
|
||
self.db.execute('DELETE FROM rooms WHERE expires_at<=?', (time.time(),))
|
||
self.db.execute('INSERT INTO rooms(room_name,user_id,expires_at,session_hash) VALUES(?,?,?,?)', (room_name, user_id, min(time.time() + ttl, session['expires_at']), hashed))
|
||
|
||
def resolve_room(self, room_name):
|
||
with self.lock:
|
||
row = self.db.execute("SELECT u.* FROM rooms r JOIN users u ON r.user_id=u.id JOIN sessions s ON s.hash=r.session_hash AND s.user_id=u.id WHERE r.room_name=? AND r.expires_at>? AND s.expires_at>? AND u.desired!='deleted' AND u.status='ready'", (room_name, time.time(), time.time())).fetchone()
|
||
return dict(row) if row else None
|