517 lines
16 KiB
Python
517 lines
16 KiB
Python
"""FastAPI dependencies (db session, storage, publisher, service-token auth)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import datetime as dt
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from typing import Annotated, Any
|
|
from uuid import UUID
|
|
|
|
from fastapi import Depends, Header, HTTPException, Request
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from ..core.api_keys import hash_api_key
|
|
from ..core.auth import AuthError, TokenExpiredError, TokenInvalidError, verify_access_token
|
|
from ..core.auth_refresh import RefreshTokenStore
|
|
from ..core.config import get_settings
|
|
from ..core.db.models import User
|
|
from ..core.db.session import create_session_factory
|
|
from ..core.logging import get_logger
|
|
from ..core.mq.publisher import Publisher
|
|
from ..core.notifications.publisher import NotificationPublisher
|
|
from ..core.rate_limit import RateLimiter
|
|
from ..core.s3.port import Storage
|
|
from ..core.security.passwords import hash_password
|
|
from ..core.tokens import hash_token
|
|
|
|
log = get_logger(__name__)
|
|
|
|
|
|
async def get_db_session() -> AsyncIterator[AsyncSession]:
|
|
factory = create_session_factory()
|
|
async with factory() as session:
|
|
yield session
|
|
|
|
|
|
AsyncSessionDep = Annotated[AsyncSession, Depends(get_db_session)]
|
|
|
|
|
|
def get_storage(request: Request) -> Storage:
|
|
storage: Storage = request.app.state.storage
|
|
return storage
|
|
|
|
|
|
StorageDep = Annotated[Storage, Depends(get_storage)]
|
|
|
|
|
|
def get_publisher(request: Request) -> Publisher:
|
|
publisher: Publisher = request.app.state.publisher
|
|
return publisher
|
|
|
|
|
|
PublisherDep = Annotated[Publisher, Depends(get_publisher)]
|
|
|
|
|
|
def get_notification_publisher(request: Request) -> NotificationPublisher:
|
|
publisher: NotificationPublisher = request.app.state.notification_publisher
|
|
return publisher
|
|
|
|
|
|
NotificationPublisherDep = Annotated[NotificationPublisher, Depends(get_notification_publisher)]
|
|
|
|
|
|
def get_redis(request: Request) -> Any:
|
|
"""Return the app-state async Redis client.
|
|
|
|
Set on app.state.redis in the api lifespan. Used by the refresh-token store.
|
|
"""
|
|
redis: Any = request.app.state.redis
|
|
return redis
|
|
|
|
|
|
RedisDep = Annotated[Any, Depends(get_redis)]
|
|
|
|
|
|
def get_refresh_store(request: Request) -> RefreshTokenStore:
|
|
"""Build a RefreshTokenStore from the app-state Redis client.
|
|
|
|
Raises 503 if Redis is unavailable — webUI auth cannot function without it.
|
|
"""
|
|
redis: Any = getattr(request.app.state, "redis", None)
|
|
if redis is None:
|
|
raise HTTPException(status_code=503, detail="refresh-token store unavailable")
|
|
settings = get_settings()
|
|
return RefreshTokenStore(redis, ttl_seconds=settings.jwt_refresh_ttl_days * 24 * 3600)
|
|
|
|
|
|
RefreshStoreDep = Annotated[RefreshTokenStore, Depends(get_refresh_store)]
|
|
|
|
|
|
async def require_service_token(
|
|
session: AsyncSessionDep,
|
|
authorization: Annotated[str | None, Header()] = None,
|
|
) -> None:
|
|
"""Validate `Authorization: Bearer <token>` against service_tokens table.
|
|
|
|
Raises 401 on missing/invalid/revoked token. Logs last_used_at on success.
|
|
"""
|
|
if not authorization or not authorization.lower().startswith("bearer "):
|
|
raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
|
|
|
|
raw = authorization[7:].strip()
|
|
token_hash = hash_token(raw)
|
|
result = await session.execute(
|
|
text("SELECT id FROM service_tokens WHERE token_hash = :h AND revoked = FALSE"),
|
|
{"h": token_hash},
|
|
)
|
|
row = result.first()
|
|
if row is None:
|
|
raise HTTPException(status_code=401, detail="Invalid or revoked token")
|
|
|
|
await session.execute(
|
|
text("UPDATE service_tokens SET last_used_at = now() WHERE id = :id"),
|
|
{"id": row[0]},
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
AuthDep = Annotated[None, Depends(require_service_token)]
|
|
|
|
|
|
async def get_or_create_user_for_telegram(
|
|
session: AsyncSession,
|
|
telegram_id: int,
|
|
*,
|
|
profile: dict[str, object] | None = None,
|
|
verified: bool = True,
|
|
) -> User:
|
|
"""Fetch or create a user identified by telegram_id.
|
|
|
|
On creation stores the profile snapshot, verification flag, and binding time.
|
|
On existing row optionally refreshes the profile snapshot.
|
|
"""
|
|
result = await session.execute(
|
|
text(
|
|
"SELECT id, telegram_id, email, password_hash, is_active, "
|
|
"created_at, credits_left, telegram_verified "
|
|
"FROM users WHERE telegram_id = :t"
|
|
),
|
|
{"t": telegram_id},
|
|
)
|
|
row = result.first()
|
|
if row:
|
|
# Refresh profile snapshot on every /start so admin sees current data.
|
|
if profile:
|
|
await session.execute(
|
|
text(
|
|
"UPDATE users SET telegram_profile_json = :p, telegram_verified = :v "
|
|
"WHERE id = :u"
|
|
),
|
|
{"p": json.dumps(profile, ensure_ascii=False), "v": verified, "u": row[0]},
|
|
)
|
|
await session.commit()
|
|
return User(
|
|
id=row[0],
|
|
telegram_id=row[1],
|
|
email=row[2],
|
|
password_hash=row[3],
|
|
is_active=row[4],
|
|
created_at=row[5],
|
|
credits_left=row[6],
|
|
telegram_verified=verified if profile else row[7],
|
|
)
|
|
|
|
bound_at = dt.datetime.now(tz=dt.UTC)
|
|
profile_json = json.dumps(profile, ensure_ascii=False) if profile else None
|
|
insert = await session.execute(
|
|
text(
|
|
"INSERT INTO users (telegram_id, credits_left, telegram_profile_json, "
|
|
"telegram_verified, telegram_bound_at) "
|
|
"VALUES (:t, 0, :p, :v, :b) "
|
|
"RETURNING id, telegram_id, created_at, credits_left, is_active, "
|
|
"telegram_verified"
|
|
),
|
|
{"t": telegram_id, "p": profile_json, "v": verified, "b": bound_at},
|
|
)
|
|
new = insert.first()
|
|
assert new is not None
|
|
await session.commit()
|
|
return User(
|
|
id=new[0],
|
|
telegram_id=new[1],
|
|
created_at=new[2],
|
|
credits_left=new[3],
|
|
is_active=new[4],
|
|
telegram_verified=new[5],
|
|
)
|
|
|
|
|
|
async def get_or_create_user_by_id(session: AsyncSession, user_id: UUID) -> User | None:
|
|
"""Fetch an existing user by UUID. Returns None if not found."""
|
|
result = await session.execute(
|
|
text("SELECT id, telegram_id, created_at, credits_left FROM users WHERE id = :u"),
|
|
{"u": user_id},
|
|
)
|
|
row = result.first()
|
|
if row is None:
|
|
return None
|
|
return User(id=row[0], telegram_id=row[1], created_at=row[2], credits_left=row[3])
|
|
|
|
|
|
async def bind_telegram_to_user(
|
|
session: AsyncSession,
|
|
user_id: UUID,
|
|
telegram_id: int,
|
|
) -> None:
|
|
"""Link a Telegram id to an existing user (web -> Telegram).
|
|
|
|
Raises HTTPException 409 if the telegram_id is already bound to another user.
|
|
"""
|
|
dup = await session.execute(
|
|
text("SELECT id FROM users WHERE telegram_id = :t AND id != :u"),
|
|
{"t": telegram_id, "u": user_id},
|
|
)
|
|
if dup.first() is not None:
|
|
raise HTTPException(status_code=409, detail="Telegram id already bound to another account")
|
|
|
|
await session.execute(
|
|
text("UPDATE users SET telegram_id = :t, telegram_bound_at = now() WHERE id = :u"),
|
|
{"t": telegram_id, "u": user_id},
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
async def set_user_password(session: AsyncSession, user_id: UUID, password: str) -> None:
|
|
"""Set/rotate a web password for a user (Telegram -> web UI access)."""
|
|
hashed = hash_password(password)
|
|
await session.execute(
|
|
text("UPDATE users SET password_hash = :p WHERE id = :u"),
|
|
{"p": hashed, "u": user_id},
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
async def fetch_user_by_email(session: AsyncSession, email: str) -> User | None:
|
|
"""Fetch a user by email (case-sensitive — normalize upstream). Returns None if not found."""
|
|
result = await session.execute(
|
|
text(
|
|
"SELECT id, telegram_id, email, name, password_hash, is_active, created_at, credits_left "
|
|
"FROM users WHERE email = :e"
|
|
),
|
|
{"e": email},
|
|
)
|
|
row = result.first()
|
|
if row is None:
|
|
return None
|
|
return User(
|
|
id=row[0],
|
|
telegram_id=row[1],
|
|
email=row[2],
|
|
name=row[3],
|
|
password_hash=row[4],
|
|
is_active=row[5],
|
|
created_at=row[6],
|
|
credits_left=row[7],
|
|
)
|
|
|
|
|
|
async def fetch_user_by_id_full(session: AsyncSession, user_id: UUID) -> User | None:
|
|
"""Fetch a user by UUID including web-auth columns."""
|
|
result = await session.execute(
|
|
text(
|
|
"SELECT id, telegram_id, email, name, password_hash, is_active, created_at, credits_left "
|
|
"FROM users WHERE id = :u"
|
|
),
|
|
{"u": user_id},
|
|
)
|
|
row = result.first()
|
|
if row is None:
|
|
return None
|
|
return User(
|
|
id=row[0],
|
|
telegram_id=row[1],
|
|
email=row[2],
|
|
name=row[3],
|
|
password_hash=row[4],
|
|
is_active=row[5],
|
|
created_at=row[6],
|
|
credits_left=row[7],
|
|
)
|
|
|
|
|
|
async def create_email_user(
|
|
session: AsyncSession, *, email: str, name: str | None, password_hash: str
|
|
) -> User:
|
|
"""Insert a new email/password user with 0 credits and return it."""
|
|
result = await session.execute(
|
|
text(
|
|
"INSERT INTO users (email, name, password_hash, credits_left) "
|
|
"VALUES (:e, :n, :p, 0) "
|
|
"RETURNING id, telegram_id, email, name, password_hash, is_active, created_at, credits_left"
|
|
),
|
|
{"e": email, "n": name, "p": password_hash},
|
|
)
|
|
row = result.first()
|
|
assert row is not None
|
|
await session.commit()
|
|
return User(
|
|
id=row[0],
|
|
telegram_id=row[1],
|
|
email=row[2],
|
|
name=row[3],
|
|
password_hash=row[4],
|
|
is_active=row[5],
|
|
created_at=row[6],
|
|
credits_left=row[7],
|
|
)
|
|
|
|
|
|
def get_rate_limiter(request: Request) -> RateLimiter:
|
|
"""Return the app-state rate limiter (Redis in prod, Memory in tests)."""
|
|
limiter: RateLimiter = request.app.state.rate_limiter
|
|
return limiter
|
|
|
|
|
|
RateLimiterDep = Annotated[RateLimiter, Depends(get_rate_limiter)]
|
|
|
|
|
|
class ApiKeyAuth:
|
|
"""Validated B2B API key + its owning user id."""
|
|
|
|
def __init__(self, api_key_id: UUID, user_id: UUID, rate_limit_rps: int) -> None:
|
|
self.api_key_id = api_key_id
|
|
self.user_id = user_id
|
|
self.rate_limit_rps = rate_limit_rps
|
|
|
|
|
|
async def _maybe_reset_monthly_quota(session: AsyncSession, api_key_id: UUID) -> None:
|
|
"""Reset monthly_used/resets_at if the quota window has expired."""
|
|
await session.execute(
|
|
text(
|
|
"UPDATE api_keys "
|
|
"SET monthly_used = 0, resets_at = now() + interval '1 month' "
|
|
"WHERE id = :k AND resets_at < now()"
|
|
),
|
|
{"k": api_key_id},
|
|
)
|
|
|
|
|
|
async def require_api_key(
|
|
session: AsyncSessionDep,
|
|
rate_limiter: RateLimiterDep,
|
|
x_api_key: Annotated[str | None, Header(alias="X-API-Key")] = None,
|
|
) -> ApiKeyAuth:
|
|
"""Validate `X-API-Key` header, rate-limit, and return key metadata.
|
|
|
|
Raises 401 on missing/invalid/revoked key. Raises 429 when rate-limited
|
|
or monthly quota exhausted.
|
|
"""
|
|
if not x_api_key:
|
|
raise HTTPException(status_code=401, detail="Missing X-API-Key header")
|
|
|
|
key_hash = hash_api_key(x_api_key)
|
|
result = await session.execute(
|
|
text(
|
|
"SELECT id, user_id, rate_limit_rps, monthly_quota, monthly_used, revoked "
|
|
"FROM api_keys "
|
|
"WHERE key_hash = :h"
|
|
),
|
|
{"h": key_hash},
|
|
)
|
|
row = result.first()
|
|
if row is None:
|
|
raise HTTPException(status_code=401, detail="Invalid API key")
|
|
|
|
api_key_id, user_id, rate_limit_rps, monthly_quota, monthly_used, revoked = row
|
|
if revoked:
|
|
raise HTTPException(status_code=401, detail="Revoked API key")
|
|
|
|
await _maybe_reset_monthly_quota(session, api_key_id)
|
|
|
|
# Check monthly quota (if configured) after potential reset.
|
|
if monthly_quota is not None and int(monthly_used) >= int(monthly_quota):
|
|
raise HTTPException(status_code=429, detail="Monthly quota exceeded")
|
|
|
|
# Apply token-bucket rate limit per key id.
|
|
limit = int(rate_limit_rps or get_settings().b2b_default_rate_limit_rps)
|
|
rl_result = await rate_limiter.allow(f"rate_limit:{api_key_id}", limit)
|
|
if not rl_result.allowed:
|
|
retry_after = max(1, int(rl_result.retry_after_sec or 1))
|
|
raise HTTPException(
|
|
status_code=429,
|
|
detail="Rate limit exceeded",
|
|
headers={"Retry-After": str(retry_after)},
|
|
)
|
|
|
|
await session.execute(
|
|
text("UPDATE api_keys SET last_used_at = now() WHERE id = :id"),
|
|
{"id": api_key_id},
|
|
)
|
|
await session.commit()
|
|
|
|
return ApiKeyAuth(api_key_id=api_key_id, user_id=user_id, rate_limit_rps=limit)
|
|
|
|
|
|
ApiKeyAuthDep = Annotated[ApiKeyAuth, Depends(require_api_key)]
|
|
|
|
|
|
async def fetch_document_status_for_user(
|
|
session: AsyncSession, document_id: UUID, user_id: UUID
|
|
) -> dict[str, object] | None:
|
|
"""Fetch document + report scoped to a specific user (B2B API)."""
|
|
result = await session.execute(
|
|
text(
|
|
"SELECT d.id, d.status, d.stage, d.filename, d.created_at, "
|
|
" r.markdown, r.content_json, r.model_used, "
|
|
" r.prompt_tokens, r.eval_tokens, r.latency_ms "
|
|
"FROM documents d "
|
|
"LEFT JOIN reports r ON r.document_id = d.id "
|
|
"WHERE d.id = :d AND d.user_id = :u"
|
|
),
|
|
{"d": document_id, "u": user_id},
|
|
)
|
|
row = result.first()
|
|
if row is None:
|
|
return None
|
|
return {
|
|
"id": row[0],
|
|
"status": row[1],
|
|
"stage": row[2],
|
|
"filename": row[3],
|
|
"created_at": row[4],
|
|
"markdown": row[5],
|
|
"content_json": row[6],
|
|
"model_used": row[7],
|
|
"prompt_tokens": row[8],
|
|
"eval_tokens": row[9],
|
|
"latency_ms": row[10],
|
|
}
|
|
|
|
|
|
class CurrentUser:
|
|
"""Authenticated user extracted from a Bearer JWT."""
|
|
|
|
def __init__(self, user_id: UUID, telegram_id: int) -> None:
|
|
self.user_id = user_id
|
|
self.telegram_id = telegram_id
|
|
|
|
|
|
async def require_current_user(
|
|
session: AsyncSessionDep,
|
|
authorization: Annotated[str | None, Header()] = None,
|
|
) -> CurrentUser:
|
|
"""Validate `Authorization: Bearer <user-jwt>` and return the user.
|
|
|
|
This is the common auth gate for bot, web, and Mini App users.
|
|
"""
|
|
if not authorization or not authorization.lower().startswith("bearer "):
|
|
raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
|
|
|
|
token = authorization[7:].strip()
|
|
try:
|
|
claims = verify_access_token(token)
|
|
except TokenExpiredError as exc:
|
|
raise HTTPException(status_code=401, detail="Token expired") from exc
|
|
except TokenInvalidError as exc:
|
|
raise HTTPException(status_code=401, detail="Invalid token") from exc
|
|
except AuthError as exc:
|
|
raise HTTPException(status_code=401, detail=str(exc)) from exc
|
|
|
|
# Ensure the user still exists (defense in depth: tokens are stateless,
|
|
# but a deleted user should not be able to use them).
|
|
result = await session.execute(
|
|
text("SELECT id FROM users WHERE id = :u"),
|
|
{"u": claims.sub},
|
|
)
|
|
if result.first() is None:
|
|
raise HTTPException(status_code=401, detail="User not found")
|
|
|
|
return CurrentUser(user_id=claims.sub, telegram_id=claims.telegram_id)
|
|
|
|
|
|
CurrentUserDep = Annotated[CurrentUser, Depends(require_current_user)]
|
|
|
|
|
|
async def fetch_document_status(
|
|
session: AsyncSession, document_id: UUID, user_id: UUID
|
|
) -> dict[str, object] | None:
|
|
result = await session.execute(
|
|
text(
|
|
"SELECT d.id, d.status, d.stage, d.filename, d.created_at, "
|
|
" r.markdown, r.content_json, r.model_used, "
|
|
" r.prompt_tokens, r.eval_tokens, r.latency_ms "
|
|
"FROM documents d "
|
|
"LEFT JOIN reports r ON r.document_id = d.id "
|
|
"WHERE d.id = :d AND d.user_id = :u"
|
|
),
|
|
{"d": document_id, "u": user_id},
|
|
)
|
|
row = result.first()
|
|
if row is None:
|
|
return None
|
|
return {
|
|
"id": row[0],
|
|
"status": row[1],
|
|
"stage": row[2],
|
|
"filename": row[3],
|
|
"created_at": row[4],
|
|
"markdown": row[5],
|
|
"content_json": row[6],
|
|
"model_used": row[7],
|
|
"prompt_tokens": row[8],
|
|
"eval_tokens": row[9],
|
|
"latency_ms": row[10],
|
|
}
|
|
|
|
|
|
async def get_credits(session: AsyncSession, user_id: UUID) -> int:
|
|
result = await session.execute(
|
|
text("SELECT credits_left FROM users WHERE id = :u"),
|
|
{"u": user_id},
|
|
)
|
|
row = result.first()
|
|
if row is None:
|
|
return 0
|
|
return int(row[0])
|