Files
simple-chat-api/users.py
T
2026-07-03 21:54:19 +08:00

160 lines
5.0 KiB
Python

from dataclasses import dataclass
from datetime import UTC, datetime
import hmac
import secrets
from pathlib import Path
from uuid import uuid4
import aiosqlite
import bcrypt
class UsernameExistsError(Exception):
pass
class InvalidCredentialsError(Exception):
pass
class UserStoreError(Exception):
pass
@dataclass(frozen=True, slots=True)
class RegisteredUser:
user_id: str
api_key: str
created_at: datetime
@dataclass(frozen=True, slots=True)
class AuthenticatedUser:
user_id: str
_SCHEMA = """
CREATE TABLE IF NOT EXISTS users (
user_id TEXT PRIMARY KEY,
username TEXT UNIQUE NOT NULL,
password_hash TEXT NOT NULL,
api_key TEXT UNIQUE NOT NULL,
created_at TEXT NOT NULL
);
"""
class UserStore:
def __init__(self, db_path: Path) -> None:
self.db_path = db_path
async def init(self) -> None:
try:
self.db_path.parent.mkdir(parents=True, exist_ok=True)
async with aiosqlite.connect(self.db_path) as db:
await db.execute("PRAGMA journal_mode=WAL;")
await db.execute("PRAGMA busy_timeout=5000;")
await db.execute("PRAGMA foreign_keys=ON;")
await db.executescript(_SCHEMA)
await db.commit()
except aiosqlite.Error as exc:
raise UserStoreError("could not initialize user store") from exc
async def register(self, username: str, password: str) -> RegisteredUser:
password_hash = bcrypt.hashpw(
password.encode("utf-8"), bcrypt.gensalt()
).decode("utf-8")
now = datetime.now(UTC)
user = RegisteredUser(
user_id=str(uuid4()),
api_key=secrets.token_urlsafe(32),
created_at=now,
)
try:
async with aiosqlite.connect(self.db_path) as db:
await db.execute(
"INSERT INTO users (user_id, username, password_hash, api_key, created_at) "
"VALUES (?, ?, ?, ?, ?)",
(
user.user_id,
username,
password_hash,
user.api_key,
now.isoformat(),
),
)
await db.commit()
except aiosqlite.IntegrityError as exc:
raise UsernameExistsError(username) from exc
except aiosqlite.Error as exc:
raise UserStoreError("could not register user") from exc
return user
async def login(self, username: str, password: str) -> RegisteredUser:
try:
async with aiosqlite.connect(self.db_path) as db:
cursor = await db.execute(
"SELECT user_id, password_hash, api_key, created_at "
"FROM users WHERE username = ?",
(username,),
)
row = await cursor.fetchone()
except aiosqlite.Error as exc:
raise UserStoreError("could not query user") from exc
if row is None:
_run_constant_time_password_check()
raise InvalidCredentialsError("invalid username or password")
user_id, password_hash, api_key, created_at = row
try:
password_ok = bcrypt.checkpw(
password.encode("utf-8"), password_hash.encode("utf-8")
)
except ValueError as exc:
raise UserStoreError("stored password hash is corrupt") from exc
if not password_ok:
raise InvalidCredentialsError("invalid username or password")
try:
parsed_created_at = datetime.fromisoformat(created_at)
except ValueError as exc:
raise UserStoreError("stored created_at is corrupt") from exc
return RegisteredUser(
user_id=user_id,
api_key=api_key,
created_at=parsed_created_at,
)
async def authenticate(self, api_key: str | None) -> AuthenticatedUser | None:
if not api_key:
return None
try:
async with aiosqlite.connect(self.db_path) as db:
cursor = await db.execute(
"SELECT api_key, user_id FROM users WHERE api_key = ?",
(api_key,),
)
row = await cursor.fetchone()
except aiosqlite.Error:
return None
if row is None:
return None
stored_key, user_id = row
if not hmac.compare_digest(api_key, stored_key):
return None
return AuthenticatedUser(user_id=user_id)
def _run_constant_time_password_check() -> None:
"""Perform a dummy bcrypt comparison to keep timing roughly constant.
Equalizes the response time between an unknown username (no hash to
compare against) and a wrong password for a known username, reducing the
feasibility of username enumeration via timing.
"""
dummy_hash = bcrypt.gensalt()
bcrypt.checkpw(b"dummy-password", dummy_hash)