reqkey 0.1.0__py3-none-any.whl

reqkey/asgi.py

@@ -0,0 +1,612 @@
+"""Generic ASGI integration for ReqKey."""
+
+from __future__ import annotations
+
+import inspect
+import logging
+import time
+import zlib
+from collections.abc import Awaitable, Callable, Iterable, Mapping
+from datetime import datetime
+from typing import Any, Literal, Protocol
+from urllib.parse import urlencode
+
+from starlette.datastructures import Headers, MutableHeaders
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from .client import DEFAULT_BASE_URL, DEFAULT_TIMEOUT_SECONDS, AsyncReqKey
+from .exceptions import ReqKeyConfigurationError, ReqKeyError
+from .models import VerificationReason, VerificationResult
+
+logger = logging.getLogger("reqkey")
+
+MAX_RESPONSE_BODY_CHARACTERS = 1000
+_MAX_CAPTURE_BYTES = MAX_RESPONSE_BODY_CHARACTERS * 4 + 4
+
+Mode = Literal["validate", "ingest", "both"]
+FailureMode = Literal["closed", "open"]
+KeyLocation = Literal["header", "query", "cookie"]
+KeyScheme = Literal["raw", "bearer"]
+CreditsResolver = Callable[[Request], int | Awaitable[int]]
+ProtectionResolver = Callable[[Request], bool | Awaitable[bool]]
+ConsumerKeyResolver = Callable[[Request], str | None | Awaitable[str | None]]
+RequestIdResolver = Callable[[Request], str | None | Awaitable[str | None]]
+PathResolver = Callable[[Request], str | Awaitable[str]]
+
+DEFAULT_ERROR_MESSAGES: dict[str, str] = {
+    "missing_api_key": "An API key is required.",
+    "invalid_api_key": "The API key is invalid or inactive.",
+    "insufficient_credits": "The API key has insufficient credits.",
+    "access_denied": "The API key is not allowed to access this API.",
+    "rate_limited": "The API key has exceeded its rate limit.",
+    "reqkey_unavailable": "API key verification is temporarily unavailable.",
+}
+
+DEFAULT_EXCLUDED_HEADERS = frozenset(
+    {
+        "authorization",
+        "cookie",
+        "proxy-authorization",
+        "set-cookie",
+        "x-api-key",
+    }
+)
+
+
+class _ResponseBodyCapture:
+    """Keep at most the bytes needed for the first 1,000 decoded characters."""
+
+    def __init__(self, *, enabled: bool, headers: Headers) -> None:
+        content_type = headers.get("content-type", "").lower()
+        self.enabled = enabled and (
+            not content_type
+            or any(
+                marker in content_type
+                for marker in (
+                    "json",
+                    "text/",
+                    "xml",
+                    "javascript",
+                    "x-www-form-urlencoded",
+                )
+            )
+        )
+        self._body = bytearray()
+        self._decompressor = (
+            zlib.decompressobj(16 + zlib.MAX_WBITS)
+            if self.enabled and headers.get("content-encoding", "").lower() == "gzip"
+            else None
+        )
+
+    def feed(self, chunk: bytes) -> None:
+        if not self.enabled or not chunk or len(self._body) >= _MAX_CAPTURE_BYTES:
+            return
+        remaining = _MAX_CAPTURE_BYTES - len(self._body)
+        if self._decompressor is None:
+            self._body.extend(chunk[:remaining])
+            return
+        try:
+            self._body.extend(self._decompressor.decompress(chunk, remaining))
+        except zlib.error:
+            self.enabled = False
+            self._body.clear()
+            logger.debug("Could not decompress a gzip response for ReqKey analytics.")
+
+    def text(self) -> str | None:
+        if not self.enabled:
+            return None
+        if self._decompressor is not None and len(self._body) < _MAX_CAPTURE_BYTES:
+            try:
+                remaining = _MAX_CAPTURE_BYTES - len(self._body)
+                self._body.extend(self._decompressor.flush(remaining))
+            except zlib.error:
+                logger.debug("Could not finish gzip response capture for ReqKey analytics.")
+        return self._body.decode("utf-8", errors="replace")[:MAX_RESPONSE_BODY_CHARACTERS]
+
+
+class AsyncReqKeyLike(Protocol):
+    async def verify(
+        self,
+        key: str,
+        *,
+        api_id: str | None = None,
+        credits: int = 1,
+        resource: str | None = None,
+    ) -> VerificationResult: ...
+
+    async def ingest(
+        self,
+        request_id: str | None = None,
+        *,
+        api_id: str | None = None,
+        method: str | None = None,
+        endpoint: str | None = None,
+        path: str | None = None,
+        status_code: int | None = None,
+        latency_ms: int | None = None,
+        client_ip: str | None = None,
+        user_agent: str | None = None,
+        user_id: str | None = None,
+        query_params: Mapping[str, Any] | None = None,
+        request_headers: Mapping[str, str] | None = None,
+        response_headers: Mapping[str, str] | None = None,
+        request_body: str | None = None,
+        response_body: str | None = None,
+        timestamp: datetime | str | None = None,
+    ) -> None: ...
+
+
+class ReqKeyMiddleware:
+    """Pure ASGI middleware for ASGI 3 applications.
+
+    `mode="validate"` gates requests only, `mode="ingest"` records traffic
+    only, and `mode="both"` validates before the handler and then performs a
+    correlated, awaited ingestion before releasing the response.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        *,
+        api_id: str,
+        project_key: str | None = None,
+        root_key: str | None = None,
+        client: AsyncReqKeyLike | None = None,
+        base_url: str = DEFAULT_BASE_URL,
+        timeout: float = DEFAULT_TIMEOUT_SECONDS,
+        mode: Mode = "both",
+        enabled: bool = True,
+        key_location: KeyLocation = "header",
+        key_name: str = "X-API-Key",
+        key_scheme: KeyScheme = "raw",
+        get_consumer_key: ConsumerKeyResolver | None = None,
+        credits: int | CreditsResolver = 1,
+        exclude_paths: Iterable[str] = (),
+        skip_methods: Iterable[str] = ("OPTIONS",),
+        should_protect: ProtectionResolver | None = None,
+        request_id_resolver: RequestIdResolver | None = None,
+        path_resolver: PathResolver | None = None,
+        error_messages: Mapping[str, str] | None = None,
+        capture_query_params: bool = False,
+        capture_request_headers: bool = False,
+        capture_response_headers: bool = False,
+        capture_response_body: bool = False,
+        capture_client_ip: bool = False,
+        capture_user_agent: bool = True,
+        excluded_headers: Iterable[str] = (),
+        failure_mode: FailureMode = "closed",
+    ) -> None:
+        if not api_id.strip():
+            raise ReqKeyConfigurationError("api_id cannot be empty.")
+        if client is not None and (project_key is not None or root_key is not None):
+            raise ReqKeyConfigurationError(
+                "Pass client or project_key/root_key, not both."
+            )
+        if mode not in {"validate", "ingest", "both"}:
+            raise ReqKeyConfigurationError(
+                "mode must be 'validate', 'ingest', or 'both'."
+            )
+        if key_location not in {"header", "query", "cookie"}:
+            raise ReqKeyConfigurationError(
+                "key_location must be 'header', 'query', or 'cookie'."
+            )
+        if key_scheme not in {"raw", "bearer"}:
+            raise ReqKeyConfigurationError("key_scheme must be 'raw' or 'bearer'.")
+        if key_location != "header" and key_scheme != "raw":
+            raise ReqKeyConfigurationError(
+                "Query-parameter and cookie keys must use the raw scheme."
+            )
+        if failure_mode not in {"closed", "open"}:
+            raise ReqKeyConfigurationError("failure_mode must be 'closed' or 'open'.")
+        if not callable(credits) and (
+            isinstance(credits, bool) or not isinstance(credits, int) or credits < 0
+        ):
+            raise ReqKeyConfigurationError(
+                "credits must be a non-negative integer or a callable."
+            )
+        self.app = app
+        self.api_id = api_id
+        self.client = client or AsyncReqKey(
+            project_key=project_key,
+            root_key=root_key,
+            base_url=base_url,
+            timeout=timeout,
+        )
+        self._owns_client = client is None
+        self.mode = mode
+        self.enabled = enabled
+        self.key_location = key_location
+        self.key_name = key_name
+        self.key_scheme = key_scheme
+        self.get_consumer_key = get_consumer_key
+        self.credits = credits
+        self.exclude_paths = tuple(exclude_paths)
+        self.skip_methods = frozenset(method.upper() for method in skip_methods)
+        self.should_protect = should_protect
+        self.request_id_resolver = request_id_resolver
+        self.path_resolver = path_resolver
+        self.error_messages = {**DEFAULT_ERROR_MESSAGES, **dict(error_messages or {})}
+        self.capture_query_params = capture_query_params
+        self.capture_request_headers = capture_request_headers
+        self.capture_response_headers = capture_response_headers
+        self.capture_response_body = capture_response_body
+        self.capture_client_ip = capture_client_ip
+        self.capture_user_agent = capture_user_agent
+        self.excluded_headers = DEFAULT_EXCLUDED_HEADERS | {
+            header.lower() for header in excluded_headers
+        } | {key_name.lower()}
+        self.failure_mode = failure_mode
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] == "lifespan":
+            try:
+                await self.app(scope, receive, send)
+            finally:
+                if self._owns_client and isinstance(self.client, AsyncReqKey):
+                    await self.client.close()
+            return
+
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        scope = dict(scope)
+        scope["state"] = dict(scope.get("state", {}))
+        request = Request(scope, receive=receive)
+        if not self.enabled or not await self._applies_to(request):
+            await self.app(scope, receive, send)
+            return
+
+        decision: VerificationResult | None = None
+        validation_time_ms: float | None = None
+
+        if self.mode in {"validate", "both"}:
+            consumer_key = await self._consumer_key(request)
+            if consumer_key is None:
+                await self._respond(
+                    scope,
+                    receive,
+                    send,
+                    status_code=401,
+                    error="missing_api_key",
+                )
+                return
+
+            validation_started = time.perf_counter()
+            try:
+                decision = await self.client.verify(
+                    consumer_key,
+                    api_id=self.api_id,
+                    credits=await self._credit_cost(request),
+                    resource=await self._resource_path(request),
+                )
+                validation_time_ms = (time.perf_counter() - validation_started) * 1000
+            except ReqKeyError as exc:
+                validation_time_ms = (time.perf_counter() - validation_started) * 1000
+                if self.failure_mode == "open":
+                    scope.setdefault("state", {})["reqkey_error"] = exc
+                else:
+                    logger.warning("ReqKey validation failed closed: %s", exc)
+                    await self._respond(
+                        scope,
+                        receive,
+                        send,
+                        status_code=503,
+                        error="reqkey_unavailable",
+                    )
+                    return
+
+            if decision is not None and not decision.valid:
+                status_code, error = self._denial(decision)
+                headers = {}
+                if decision.retry_after is not None:
+                    headers["Retry-After"] = str(max(0, int(decision.retry_after)))
+                await self._respond(
+                    scope,
+                    receive,
+                    send,
+                    status_code=status_code,
+                    error=error,
+                    headers=headers,
+                )
+                return
+
+            if decision is not None:
+                state = scope.setdefault("state", {})
+                state["reqkey"] = decision
+                state["reqkey_request_id"] = decision.request_id
+
+        if self.mode == "validate":
+            await self._run_validation_only(
+                scope,
+                receive,
+                send,
+                decision=decision,
+                validation_time_ms=validation_time_ms,
+            )
+            return
+
+        await self._run_with_blocking_ingest(
+            request,
+            scope,
+            receive,
+            send,
+            decision=decision,
+            validation_time_ms=validation_time_ms,
+        )
+
+    async def _run_validation_only(
+        self,
+        scope: Scope,
+        receive: Receive,
+        send: Send,
+        *,
+        decision: VerificationResult | None,
+        validation_time_ms: float | None,
+    ) -> None:
+        async def send_with_headers(message: Message) -> None:
+            self._add_decision_headers(message, decision, validation_time_ms)
+            await send(message)
+
+        await self.app(scope, receive, send_with_headers)
+
+    async def _run_with_blocking_ingest(
+        self,
+        request: Request,
+        scope: Scope,
+        receive: Receive,
+        send: Send,
+        *,
+        decision: VerificationResult | None,
+        validation_time_ms: float | None,
+    ) -> None:
+        response_status = 500
+        response_headers = Headers()
+        body_capture = _ResponseBodyCapture(enabled=False, headers=response_headers)
+        ingestion_completed = False
+        started = time.perf_counter()
+
+        async def stream_response(message: Message) -> None:
+            nonlocal response_status, response_headers, body_capture, ingestion_completed
+            if message["type"] == "http.response.start":
+                response_status = message["status"]
+                response_headers = Headers(raw=list(message.get("headers", [])))
+                body_capture = _ResponseBodyCapture(
+                    enabled=self.capture_response_body,
+                    headers=response_headers,
+                )
+                self._add_decision_headers(message, decision, validation_time_ms)
+                await send(message)
+                return
+
+            if message["type"] != "http.response.body":
+                await send(message)
+                return
+
+            body_capture.feed(message.get("body", b""))
+            if message.get("more_body", False):
+                await send(message)
+                return
+
+            await self._ingest_safely(
+                request=request,
+                decision=decision,
+                response_status=response_status,
+                latency_ms=round((time.perf_counter() - started) * 1000),
+                response_headers=response_headers,
+                response_body=body_capture.text(),
+            )
+            ingestion_completed = True
+            await send(message)
+
+        try:
+            await self.app(scope, receive, stream_response)
+        except Exception:
+            if not ingestion_completed:
+                await self._ingest_safely(
+                    request=request,
+                    decision=decision,
+                    response_status=500,
+                    latency_ms=round((time.perf_counter() - started) * 1000),
+                    response_headers=response_headers,
+                    response_body=body_capture.text(),
+                )
+            raise
+
+        if not ingestion_completed:
+            await self._ingest_safely(
+                request=request,
+                decision=decision,
+                response_status=response_status,
+                latency_ms=round((time.perf_counter() - started) * 1000),
+                response_headers=response_headers,
+                response_body=body_capture.text(),
+            )
+
+    async def _applies_to(self, request: Request) -> bool:
+        if request.method.upper() in self.skip_methods:
+            return False
+        if any(self._path_matches(request.url.path, pattern) for pattern in self.exclude_paths):
+            return False
+        if self.should_protect is None:
+            return True
+        result = self.should_protect(request)
+        return bool(await result) if inspect.isawaitable(result) else bool(result)
+
+    @staticmethod
+    def _path_matches(path: str, pattern: str) -> bool:
+        if pattern.endswith("*"):
+            return path.startswith(pattern[:-1])
+        return path == pattern
+
+    async def _consumer_key(self, request: Request) -> str | None:
+        if self.get_consumer_key is not None:
+            raw = self.get_consumer_key(request)
+            value = await raw if inspect.isawaitable(raw) else raw
+        elif self.key_location == "header":
+            value = request.headers.get(self.key_name)
+        elif self.key_location == "query":
+            value = request.query_params.get(self.key_name)
+        else:
+            value = request.cookies.get(self.key_name)
+
+        if value is None or not value.strip():
+            return None
+        value = value.strip()
+        if self.key_scheme == "bearer":
+            scheme, separator, credential = value.partition(" ")
+            if not separator or scheme.lower() != "bearer" or not credential.strip():
+                return None
+            return credential.strip()
+        return value
+
+    async def _credit_cost(self, request: Request) -> int:
+        raw = self.credits(request) if callable(self.credits) else self.credits
+        value = await raw if inspect.isawaitable(raw) else raw
+        if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+            raise ReqKeyConfigurationError(
+                "The credits resolver must return a non-negative integer."
+            )
+        return value
+
+    async def _resource_path(self, request: Request) -> str:
+        if self.path_resolver is None:
+            return request.url.path
+        raw = self.path_resolver(request)
+        value = await raw if inspect.isawaitable(raw) else raw
+        if not value:
+            raise ReqKeyConfigurationError("The path resolver returned an empty path.")
+        return value
+
+    async def _resolved_request_id(
+        self,
+        request: Request,
+        decision: VerificationResult | None,
+    ) -> str | None:
+        if decision is not None:
+            return decision.request_id
+        if self.request_id_resolver is not None:
+            raw = self.request_id_resolver(request)
+            value = await raw if inspect.isawaitable(raw) else raw
+            return value or None
+        state_request_id = getattr(request.state, "reqkey_request_id", None)
+        return state_request_id if isinstance(state_request_id, str) else None
+
+    async def _ingest_safely(
+        self,
+        *,
+        request: Request,
+        decision: VerificationResult | None,
+        response_status: int,
+        latency_ms: int,
+        response_headers: Headers,
+        response_body: str | None,
+    ) -> None:
+        resource_path = await self._resource_path(request)
+        query_params, query = self._captured_query(request)
+        full_path = resource_path + (f"?{query}" if query else "")
+        request_headers = (
+            self._filtered_headers(request.headers)
+            if self.capture_request_headers
+            else None
+        )
+        captured_response_headers = (
+            self._filtered_headers(response_headers)
+            if self.capture_response_headers
+            else None
+        )
+        client_ip = (
+            request.client.host
+            if self.capture_client_ip and request.client is not None
+            else None
+        )
+        user_agent = request.headers.get("user-agent") if self.capture_user_agent else None
+
+        try:
+            await self.client.ingest(
+                await self._resolved_request_id(request, decision),
+                api_id=self.api_id,
+                method=request.method,
+                endpoint=resource_path,
+                path=full_path,
+                status_code=response_status,
+                latency_ms=latency_ms,
+                client_ip=client_ip,
+                user_agent=user_agent,
+                query_params=query_params,
+                request_headers=request_headers,
+                response_headers=captured_response_headers,
+                response_body=response_body,
+            )
+        except ReqKeyError as exc:
+            logger.warning("ReqKey analytics ingestion failed: %s", exc)
+
+    def _captured_query(self, request: Request) -> tuple[dict[str, str] | None, str]:
+        if not self.capture_query_params:
+            return None, ""
+        pairs = [
+            (key, value)
+            for key, value in request.query_params.multi_items()
+            if not (self.key_location == "query" and key == self.key_name)
+        ]
+        return dict(pairs), urlencode(pairs)
+
+    def _filtered_headers(self, headers: Mapping[str, str]) -> dict[str, str]:
+        return {
+            key: value
+            for key, value in headers.items()
+            if key.lower() not in self.excluded_headers
+        }
+
+    @staticmethod
+    def _add_decision_headers(
+        message: Message,
+        decision: VerificationResult | None,
+        validation_time_ms: float | None,
+    ) -> None:
+        if message["type"] != "http.response.start":
+            return
+        headers = MutableHeaders(scope=message)
+        if decision is not None:
+            if decision.request_id:
+                headers["X-ReqKey-Request-ID"] = decision.request_id
+            if decision.credits_limit is not None:
+                headers["X-ReqKey-Credits-Limit"] = str(decision.credits_limit)
+            if decision.credits_remaining is not None:
+                headers["X-ReqKey-Credits-Remaining"] = str(decision.credits_remaining)
+        if validation_time_ms is not None:
+            headers["X-ReqKey-Validation-Time-Ms"] = f"{validation_time_ms:.3f}"
+
+    @staticmethod
+    def _denial(decision: VerificationResult) -> tuple[int, str]:
+        if decision.reason is VerificationReason.INSUFFICIENT_CREDITS:
+            return 402, "insufficient_credits"
+        if decision.reason is VerificationReason.RATE_LIMITED:
+            return 429, "rate_limited"
+        if decision.reason is VerificationReason.FORBIDDEN:
+            return 403, "access_denied"
+        return 401, "invalid_api_key"
+
+    async def _respond(
+        self,
+        scope: Scope,
+        receive: Receive,
+        send: Send,
+        *,
+        status_code: int,
+        error: str,
+        headers: Mapping[str, str] | None = None,
+    ) -> None:
+        response = JSONResponse(
+            {"error": error, "message": self.error_messages[error]},
+            status_code=status_code,
+            headers=dict(headers or {}),
+        )
+        await response(scope, receive, send)
+
+
+ReqKeyASGIMiddleware = ReqKeyMiddleware
+
+__all__ = ["ReqKeyASGIMiddleware", "ReqKeyMiddleware"]