464 lines
14 KiB
Python
464 lines
14 KiB
Python
"""Authentication, token, TOTP, and account bootstrap services."""
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import logging
|
|
import re
|
|
import secrets
|
|
import uuid
|
|
from datetime import datetime, timezone, timedelta
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
import jwt
|
|
import pyotp
|
|
from pwdlib import PasswordHash
|
|
from pwdlib.hashers.argon2 import Argon2Hasher
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from config import settings
|
|
from db.models import BackupCode, Quota, RefreshToken, User
|
|
|
|
_PASSWORD_DETAIL = (
|
|
"Password must be at least 12 characters and include uppercase, "
|
|
"lowercase, a number, and a special character."
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_pwd = PasswordHash([Argon2Hasher()])
|
|
|
|
|
|
def hash_password(plain: str) -> str:
|
|
"""Return the Argon2 hash of *plain*."""
|
|
return _pwd.hash(plain)
|
|
|
|
|
|
def verify_password(plain: str, hashed: str) -> bool:
|
|
"""Return True if *plain* matches *hashed* (constant-time, SEC-06).
|
|
|
|
pwdlib.verify returns True/False and internally uses constant-time comparison.
|
|
"""
|
|
try:
|
|
return _pwd.verify(plain, hashed)
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def validate_password_strength(password: str) -> None:
|
|
"""Raise ValueError with a descriptive message if *password* fails any strength rule.
|
|
|
|
Rules (AUTH-01): min 12 chars, uppercase, lowercase, digit, special char.
|
|
Callers at the API boundary should catch ValueError and map it to HTTP 422.
|
|
"""
|
|
if (
|
|
len(password) < 12
|
|
or not re.search(r"[A-Z]", password)
|
|
or not re.search(r"[a-z]", password)
|
|
or not re.search(r"[0-9]", password)
|
|
or not re.search(r"[^A-Za-z0-9]", password)
|
|
):
|
|
raise ValueError(_PASSWORD_DETAIL)
|
|
|
|
|
|
def _compute_fgp(user_agent: str, accept_lang: str) -> str:
|
|
"""Return a short fingerprint binding a token to its client context."""
|
|
return hmac.new(
|
|
settings.secret_key.encode(),
|
|
(user_agent + "\x00" + accept_lang).encode(),
|
|
hashlib.sha256,
|
|
).hexdigest()[:16]
|
|
|
|
|
|
def _jwt_private_key() -> str:
|
|
return base64.b64decode(settings.jwt_private_key).decode()
|
|
|
|
|
|
def _jwt_public_key() -> str:
|
|
return base64.b64decode(settings.jwt_public_key).decode()
|
|
|
|
|
|
def _encode_jwt(payload: dict) -> str:
|
|
return jwt.encode(payload, _jwt_private_key(), algorithm="ES256")
|
|
|
|
|
|
def _decode_jwt(
|
|
token: str,
|
|
expected_type: str,
|
|
expired_message: str,
|
|
invalid_prefix: str,
|
|
) -> dict:
|
|
try:
|
|
payload = jwt.decode(token, _jwt_public_key(), algorithms=["ES256"])
|
|
except jwt.ExpiredSignatureError as exc:
|
|
raise ValueError(expired_message) from exc
|
|
except jwt.PyJWTError as exc:
|
|
raise ValueError(f"{invalid_prefix}: {exc}") from exc
|
|
|
|
if payload.get("typ") != expected_type:
|
|
raise ValueError(f"Token type mismatch: expected '{expected_type}'")
|
|
return payload
|
|
|
|
|
|
def create_access_token(
|
|
user_id: str,
|
|
role: str,
|
|
user_agent: str = "",
|
|
accept_lang: str = "",
|
|
) -> str:
|
|
"""Return a signed JWT access token.
|
|
|
|
Claims: sub=user_id, role=role, typ='access', exp=now+access_token_expire_minutes,
|
|
fgp=fingerprint computed from client User-Agent + Accept-Language headers (D-05).
|
|
"""
|
|
now = datetime.now(timezone.utc)
|
|
payload = {
|
|
"sub": str(user_id),
|
|
"role": role,
|
|
"typ": "access",
|
|
"iat": now,
|
|
"exp": now + timedelta(minutes=settings.access_token_expire_minutes),
|
|
"jti": str(uuid.uuid4()),
|
|
"fgp": _compute_fgp(user_agent, accept_lang),
|
|
}
|
|
return _encode_jwt(payload)
|
|
|
|
|
|
def decode_access_token(token: str) -> dict:
|
|
"""Decode and validate an access token; raises ValueError on any failure.
|
|
|
|
Verifies: signature, expiry, and typ='access' (T-02-01 — prevents password-reset
|
|
tokens from being used as access tokens).
|
|
"""
|
|
return _decode_jwt(token, "access", "Token has expired", "Invalid token")
|
|
|
|
|
|
def create_password_reset_token(user_id: str) -> str:
|
|
"""Return a short-lived signed JWT for password reset.
|
|
|
|
Claims: sub=user_id, typ='password-reset', exp=now+3600s.
|
|
"""
|
|
now = datetime.now(timezone.utc)
|
|
payload = {
|
|
"sub": str(user_id),
|
|
"typ": "password-reset",
|
|
"iat": now,
|
|
"exp": now + timedelta(seconds=3600),
|
|
}
|
|
return _encode_jwt(payload)
|
|
|
|
|
|
def decode_password_reset_token(token: str) -> str:
|
|
"""Decode a password-reset token; raises ValueError on failure.
|
|
|
|
Returns the user_id string.
|
|
"""
|
|
payload = _decode_jwt(
|
|
token,
|
|
"password-reset",
|
|
"Reset token has expired",
|
|
"Invalid reset token",
|
|
)
|
|
return payload["sub"]
|
|
|
|
|
|
async def create_refresh_token(
|
|
session: AsyncSession, user_id: uuid.UUID, remember_me: bool = False
|
|
) -> str:
|
|
"""Insert a new RefreshToken row and return the raw (unhashed) token string.
|
|
|
|
The raw token is returned to the caller and set as an httpOnly cookie.
|
|
Only the SHA-256 hash is stored in the database.
|
|
|
|
remember_me=False uses the short-session TTL; remember_me=True uses the
|
|
extended-session TTL.
|
|
"""
|
|
raw = secrets.token_urlsafe(32)
|
|
token_hash = hashlib.sha256(raw.encode()).hexdigest()
|
|
now = datetime.now(timezone.utc)
|
|
ttl = (
|
|
timedelta(days=settings.refresh_token_expire_days)
|
|
if remember_me
|
|
else timedelta(hours=settings.refresh_token_expire_hours)
|
|
)
|
|
row = RefreshToken(
|
|
id=uuid.uuid4(),
|
|
user_id=user_id,
|
|
token_hash=token_hash,
|
|
expires_at=now + ttl,
|
|
revoked=False,
|
|
)
|
|
session.add(row)
|
|
await session.commit()
|
|
return raw
|
|
|
|
|
|
async def rotate_refresh_token(
|
|
session: AsyncSession, raw_token: str
|
|
) -> tuple[str, str]:
|
|
"""Rotate a refresh token: revoke the old one and issue a new one.
|
|
|
|
Returns (new_raw_token, user_id_str).
|
|
|
|
Family revocation (AUTH-07 / RFC 9700):
|
|
If the presented token has revoked=True, ALL tokens for that user are
|
|
revoked and send_security_alert_email.delay is enqueued. Then raises
|
|
ValueError("token_family_revoked").
|
|
|
|
Raises ValueError for any invalid or expired token.
|
|
"""
|
|
token_hash = hashlib.sha256(raw_token.encode()).hexdigest()
|
|
result = await session.execute(
|
|
select(RefreshToken).where(RefreshToken.token_hash == token_hash)
|
|
)
|
|
row: Optional[RefreshToken] = result.scalar_one_or_none()
|
|
|
|
if row is None:
|
|
raise ValueError("Refresh token not found")
|
|
|
|
if row.expires_at.replace(tzinfo=timezone.utc) < datetime.now(timezone.utc):
|
|
raise ValueError("Refresh token has expired")
|
|
|
|
if row.revoked:
|
|
await revoke_all_refresh_tokens(session, row.user_id)
|
|
from tasks.email_tasks import send_security_alert_email # noqa: PLC0415
|
|
send_security_alert_email.delay(str(row.user_id))
|
|
raise ValueError("token_family_revoked")
|
|
|
|
# Valid token: revoke the old one
|
|
row.revoked = True
|
|
await session.flush()
|
|
|
|
# Issue a new token in the same family
|
|
new_raw = await create_refresh_token(session, row.user_id)
|
|
return new_raw, str(row.user_id)
|
|
|
|
|
|
async def revoke_all_refresh_tokens(
|
|
session: AsyncSession, user_id: uuid.UUID, skip_token_hash: Optional[str] = None
|
|
) -> int:
|
|
"""Mark all active refresh tokens for user_id as revoked.
|
|
|
|
Returns the count of revoked tokens (supports sign-out-all-devices).
|
|
skip_token_hash: if set, the token with this hash is excluded (keep current session alive).
|
|
"""
|
|
conditions = [
|
|
RefreshToken.user_id == user_id,
|
|
RefreshToken.revoked.is_(False),
|
|
]
|
|
if skip_token_hash is not None:
|
|
conditions.append(RefreshToken.token_hash != skip_token_hash)
|
|
result = await session.execute(
|
|
select(RefreshToken).where(*conditions)
|
|
)
|
|
rows = result.scalars().all()
|
|
count = 0
|
|
for row in rows:
|
|
row.revoked = True
|
|
count += 1
|
|
await session.flush()
|
|
return count
|
|
|
|
|
|
async def provision_totp(
|
|
session: AsyncSession, user_id: uuid.UUID
|
|
) -> tuple[str, str]:
|
|
"""Generate a new TOTP secret for the user and store it (not yet enabled).
|
|
|
|
Returns (secret, provisioning_uri).
|
|
The secret is base32-encoded. provisioning_uri is suitable for QR code generation.
|
|
"""
|
|
user = await session.get(User, user_id)
|
|
if user is None:
|
|
raise ValueError("User not found")
|
|
|
|
secret = pyotp.random_base32()
|
|
user.totp_secret = secret
|
|
await session.commit()
|
|
|
|
totp = pyotp.totp.TOTP(secret)
|
|
uri = totp.provisioning_uri(user.email, issuer_name="DocuVault")
|
|
return secret, uri
|
|
|
|
|
|
async def verify_totp(
|
|
session: AsyncSession,
|
|
user_id: uuid.UUID,
|
|
code: str,
|
|
redis_client,
|
|
) -> bool:
|
|
"""Verify a TOTP code with replay prevention (AUTH-08).
|
|
|
|
valid_window=1 allows ±30s clock drift (per STATE.md recommendation).
|
|
Replay prevention: stores used codes in Redis with key 'totp_used:{user_id}:{code}'
|
|
and TTL=90s (covers the full ±30s validity window).
|
|
|
|
Returns False if the code is invalid, already used, or user has no TOTP secret.
|
|
"""
|
|
user = await session.get(User, user_id)
|
|
if user is None or not user.totp_secret:
|
|
return False
|
|
|
|
totp = pyotp.TOTP(user.totp_secret)
|
|
|
|
# Check replay prevention before verifying
|
|
replay_key = f"totp_used:{user_id}:{code}"
|
|
if await redis_client.get(replay_key):
|
|
return False # Code already used within the validity window
|
|
|
|
if not totp.verify(code, valid_window=1):
|
|
return False
|
|
|
|
# Mark as used for 90s (covers valid_window=1: ±30s = 90s total)
|
|
await redis_client.set(replay_key, "1", ex=90)
|
|
return True
|
|
|
|
|
|
def generate_backup_codes(n: int = 10) -> list[str]:
|
|
"""Return *n* random 8-character uppercase alphanumeric backup codes."""
|
|
codes = []
|
|
for _ in range(n):
|
|
code = secrets.token_hex(4).upper()
|
|
codes.append(code)
|
|
return codes
|
|
|
|
|
|
async def store_backup_codes(
|
|
session: AsyncSession, user_id: uuid.UUID, codes: list[str]
|
|
) -> None:
|
|
"""Store backup codes for a user, replacing any existing unused codes.
|
|
|
|
Each code is stored as an Argon2 hash (never plaintext, T-02-03).
|
|
Existing unused codes are deleted first to prevent accumulation.
|
|
"""
|
|
result = await session.execute(
|
|
select(BackupCode).where(
|
|
BackupCode.user_id == user_id,
|
|
BackupCode.used_at.is_(None),
|
|
)
|
|
)
|
|
for row in result.scalars().all():
|
|
await session.delete(row)
|
|
await session.flush()
|
|
|
|
for code in codes:
|
|
row = BackupCode(
|
|
id=uuid.uuid4(),
|
|
user_id=user_id,
|
|
code_hash=hash_password(code),
|
|
used_at=None,
|
|
)
|
|
session.add(row)
|
|
|
|
await session.commit()
|
|
|
|
|
|
async def verify_backup_code(
|
|
session: AsyncSession, user_id: uuid.UUID, code: str
|
|
) -> bool:
|
|
"""Verify a backup code using constant-time comparison.
|
|
|
|
Always iterates ALL unused codes (prevents timing-based enumeration, SEC-06).
|
|
On match: sets used_at=now() and commits. Returns True on first match.
|
|
Returns False if no code matches.
|
|
"""
|
|
result = await session.execute(
|
|
select(BackupCode).where(
|
|
BackupCode.user_id == user_id,
|
|
BackupCode.used_at.is_(None),
|
|
)
|
|
)
|
|
rows = result.scalars().all()
|
|
|
|
matched_row: Optional[BackupCode] = None
|
|
for row in rows:
|
|
if verify_password(code, row.code_hash):
|
|
matched_row = row # keep iterating for constant-time behavior
|
|
|
|
if matched_row is None:
|
|
return False
|
|
|
|
matched_row.used_at = datetime.now(timezone.utc)
|
|
await session.commit()
|
|
return True
|
|
|
|
|
|
async def check_hibp(password: str) -> bool:
|
|
"""Check if password appears in HaveIBeenPwned using the k-anonymity model.
|
|
|
|
Sends only the first 5 chars of the SHA-1 hash to the HIBP API (T-02-05).
|
|
Returns True if the password has been breached, False otherwise.
|
|
On network error: logs a warning and returns False (fail-open, T-02-06).
|
|
"""
|
|
sha1 = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() # noqa: S324 — HIBP k-anonymity protocol mandates SHA-1; not used as a security primitive
|
|
prefix, suffix = sha1[:5], sha1[5:]
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=5.0) as client:
|
|
resp = await client.get(
|
|
f"https://api.pwnedpasswords.com/range/{prefix}",
|
|
headers={"Add-Padding": "true"},
|
|
)
|
|
resp.raise_for_status()
|
|
except Exception as exc:
|
|
logger.warning("HIBP check failed (fail-open): %s", exc)
|
|
return False
|
|
|
|
for line in resp.text.splitlines():
|
|
parts = line.split(":")
|
|
if len(parts) == 2:
|
|
candidate_suffix, count_str = parts
|
|
if hmac.compare_digest(candidate_suffix.upper(), suffix):
|
|
try:
|
|
count = int(count_str.strip())
|
|
except ValueError:
|
|
count = 1
|
|
if count > 0:
|
|
return True
|
|
return False
|
|
|
|
|
|
async def bootstrap_admin(session: AsyncSession) -> None:
|
|
"""Idempotent admin account bootstrap.
|
|
|
|
If the users table is empty AND settings.admin_email and settings.admin_password
|
|
are both non-empty, creates an admin User row with a Quota row.
|
|
|
|
Logs a warning if env vars are missing but never raises.
|
|
"""
|
|
if not settings.admin_email or not settings.admin_password:
|
|
logger.warning(
|
|
"Admin bootstrap skipped: ADMIN_EMAIL and/or ADMIN_PASSWORD not set. "
|
|
"Set both env vars to seed the first admin account on startup."
|
|
)
|
|
return
|
|
|
|
# Check if any users exist
|
|
result = await session.execute(select(User).limit(1))
|
|
if result.scalar_one_or_none() is not None:
|
|
return
|
|
|
|
admin_id = uuid.uuid4()
|
|
admin_user = User(
|
|
id=admin_id,
|
|
handle="admin",
|
|
email=settings.admin_email,
|
|
password_hash=hash_password(settings.admin_password),
|
|
role="admin",
|
|
is_active=True,
|
|
password_must_change=False,
|
|
)
|
|
quota = Quota(
|
|
user_id=admin_id,
|
|
limit_bytes=104857600,
|
|
used_bytes=0,
|
|
)
|
|
session.add(admin_user)
|
|
await session.flush() # persist User first so Quota FK is satisfied
|
|
session.add(quota)
|
|
await session.commit()
|
|
logger.info("Admin account bootstrapped for %s", settings.admin_email)
|