spb/company-atlas

1"""Shared API helpers: JSON response class (orjson), pagination, admin/owner auth dependencies, tiny in-process TTL cache."""2from __future__ import annotations34import hashlib5import hmac6import time7from collections import OrderedDict8from typing import Annotated, Any910import orjson11from fastapi import Depends, Header, HTTPException, Query, Request, Response12from fastapi.responses import JSONResponse1314from companyatlas.config import settings151617class AtlasJSONResponse(JSONResponse):18    media_type = "application/json"1920    def render(self, content: Any) -> bytes:21        return orjson.dumps(content, option=orjson.OPT_NON_STR_KEYS | orjson.OPT_UTC_Z | orjson.OPT_SERIALIZE_NUMPY, default=_default)222324def _default(obj: Any) -> Any:25    if hasattr(obj, "isoformat"):26        return obj.isoformat()27    if hasattr(obj, "__float__"):28        return float(obj)29    if isinstance(obj, set | frozenset):30        return sorted(obj)31    raise TypeError(f"not serialisable: {type(obj)!r}")323334# ------------------------------------------------------------------------------------------------ pagination353637class PageParams:38    def __init__(self, page: int = Query(1, ge=1), per_page: int = Query(25, ge=1, le=200)):39        self.page = page40        self.per_page = per_page4142    @property43    def offset(self) -> int:44        return (self.page - 1) * self.per_page454647PageDep = Annotated[PageParams, Depends()]     # `p: PageDep` in route signatures (avoids call-in-default lint)484950def page_payload(items: list[Any], total: int, p: PageParams) -> dict[str, Any]:51    return {"items": items, "page": p.page, "per_page": p.per_page, "total": total, "pages": max(1, -(-total // p.per_page))}525354# ------------------------------------------------------------------------------------------------ auth555657def require_admin(x_ca_admin_token: str | None = Header(None, alias="X-CA-Admin-Token")) -> None:58    if not settings.admin_token or not x_ca_admin_token or not hmac.compare_digest(x_ca_admin_token, settings.admin_token):59        raise HTTPException(status_code=401, detail="admin token required")606162def owner_hash(x_ca_owner_token: str | None = Header(None, alias="X-CA-Owner-Token")) -> str:63    """Anonymous owner identity for watchlists/alerts: a client-generated random token (≥ 24 chars), stored hashed."""64    if not x_ca_owner_token or len(x_ca_owner_token) < 24 or len(x_ca_owner_token) > 200:65        raise HTTPException(status_code=401, detail="owner token required (X-CA-Owner-Token, ≥ 24 characters)")66    return hashlib.sha256(x_ca_owner_token.encode()).hexdigest()676869def client_ip(request: Request) -> str:70    fwd = request.headers.get("x-forwarded-for")71    if fwd:72        return fwd.split(",")[0].strip()73    return request.client.host if request.client else "0.0.0.0"747576# ------------------------------------------------------------------------------------------------ cache (per process)777879class TTLCache:80    def __init__(self, max_items: int = 4096):81        self._data: OrderedDict[str, tuple[float, Any]] = OrderedDict()82        self._max = max_items8384    def get(self, key: str) -> Any | None:85        item = self._data.get(key)86        if item is None:87            return None88        exp, value = item89        if exp < time.monotonic():90            self._data.pop(key, None)91            return None92        self._data.move_to_end(key)93        return value9495    def set(self, key: str, value: Any, ttl_s: float) -> None:96        self._data[key] = (time.monotonic() + ttl_s, value)97        self._data.move_to_end(key)98        while len(self._data) > self._max:99            self._data.popitem(last=False)100101    def clear(self, prefix: str | None = None) -> None:102        if prefix is None:103            self._data.clear()104        else:105            for k in [k for k in self._data if k.startswith(prefix)]:106                self._data.pop(k, None)107108109cache = TTLCache()110111112async def cached(key: str, ttl_s: float, producer):  # type: ignore[no-untyped-def]113    hit = cache.get(key)114    if hit is not None:115        return hit116    value = await producer()117    cache.set(key, value, ttl_s)118    return value119120121# ------------------------------------------------------------------------------------------------ cache headers / conditional GET122123NO_STORE = "no-store"124125126def public_cache_value(max_age_s: int) -> str:127    return f"public, max-age={max_age_s}, stale-while-revalidate={max_age_s * 2}"128129130def cached_response(request: Request, payload: Any, max_age_s: int) -> Any:131    """JSON response for a cached public aggregate: `Cache-Control: public, max-age=…` + weak ETag, `304` when `If-None-Match` matches."""132    body = AtlasJSONResponse(payload).body133    etag = 'W/"' + hashlib.sha1(body).hexdigest()[:20] + '"'134    headers = {"cache-control": public_cache_value(max_age_s), "etag": etag, "vary": "Accept-Encoding"}135    inm = request.headers.get("if-none-match")136    if inm and etag in [x.strip() for x in inm.split(",")]:137        return Response(status_code=304, headers=headers)138    return Response(content=body, media_type=AtlasJSONResponse.media_type, headers=headers)139140141def is_admin_request(request: Request) -> bool:142    tok = request.headers.get("x-ca-admin-token")143    return bool(settings.admin_token and tok and hmac.compare_digest(tok, settings.admin_token))144145146__all__ = ["NO_STORE", "AtlasJSONResponse", "PageDep", "PageParams", "TTLCache", "cache", "cached", "cached_response", "client_ip", "is_admin_request",147           "owner_hash", "page_payload", "public_cache_value", "require_admin"]148