|
| 1 | +from typing import Callable |
| 2 | + |
| 3 | +from fastapi import APIRouter, Depends, HTTPException, Request |
| 4 | +from fastapi.types import DecoratedCallable |
| 5 | + |
| 6 | +from webhook_utils.crypto import compare_sha256_signature |
| 7 | + |
| 8 | + |
| 9 | +class WebhookRouter: |
| 10 | + """Wraps a FastAPI router and adds a signature verification step for webhooks. |
| 11 | +
|
| 12 | + Example |
| 13 | + ------- |
| 14 | + >>> from fastapi import FastAPI, APIRouter |
| 15 | + >>> app = FastAPI() |
| 16 | + >>> router = APIRouter(prefix="/webhooks") |
| 17 | + >>> webhook_router = WebhookRouter(router, webhook_key="secret") |
| 18 | + >>> @webhook_router.on("/demo-webhook") |
| 19 | + ... def demo_event_handler(request: Request): |
| 20 | + ... return {"status": "ok"} |
| 21 | + >>> app.include_router(webhook_router.api_router) |
| 22 | + """ |
| 23 | + |
| 24 | + def __init__( |
| 25 | + self, |
| 26 | + api_router: APIRouter, |
| 27 | + *, |
| 28 | + webhook_key: str, |
| 29 | + header_name: str = "X-Webhook-Signature", |
| 30 | + compare_signature_fn: Callable[ |
| 31 | + [bytes, bytes, str], bool |
| 32 | + ] = compare_sha256_signature, |
| 33 | + ) -> None: |
| 34 | + """Initialize the WebhookRouter. |
| 35 | +
|
| 36 | + Parameters |
| 37 | + ---------- |
| 38 | + api_router: APIRouter |
| 39 | + The FastAPI router to wrap. |
| 40 | + webhook_key: str |
| 41 | + The webhook key to use for signature verification. |
| 42 | + header_name: str |
| 43 | + The name of the header to use for signature verification. |
| 44 | + Defaults to "X-Webhook-Signature". |
| 45 | + compare_signature_fn: Callable[[bytes, bytes, str], bool] |
| 46 | + The function to use for signature verification. |
| 47 | + Can be overridden for different signature algorithms. |
| 48 | + Defaults to `compare_sha256_signature`. |
| 49 | + """ |
| 50 | + |
| 51 | + self._api_router = api_router |
| 52 | + self._webhook_key = webhook_key.encode("utf-8") |
| 53 | + self._header_name = header_name |
| 54 | + self._compare_signature_fn = compare_signature_fn |
| 55 | + |
| 56 | + self._api_router.dependencies.append(Depends(self._validate_webhook_signature)) |
| 57 | + |
| 58 | + @property |
| 59 | + def api_router(self) -> APIRouter: |
| 60 | + """Returns the internal FastAPI router.""" |
| 61 | + return self._api_router |
| 62 | + |
| 63 | + def on(self, *args, **kwargs) -> Callable[[DecoratedCallable], DecoratedCallable]: |
| 64 | + """Adds an API route to the internal api router for this webhook event. |
| 65 | +
|
| 66 | + Parameters are the same as the `APIRouter.api_route` decorator. |
| 67 | + The default HTTP method is `POST` and is generally recommended to use. |
| 68 | + """ |
| 69 | + |
| 70 | + kwargs.setdefault("methods", ["POST"]) |
| 71 | + |
| 72 | + def decorator(fn: DecoratedCallable) -> DecoratedCallable: |
| 73 | + return self._api_router.api_route(*args, **kwargs)(fn) |
| 74 | + |
| 75 | + return decorator |
| 76 | + |
| 77 | + async def _validate_webhook_signature(self, request: Request) -> None: |
| 78 | + """Check if the webhook signature is valid. |
| 79 | +
|
| 80 | + This method is automatically added as a dependency to all routes |
| 81 | + created with `WebhookRouter.on`. |
| 82 | +
|
| 83 | + Raises HTTPException (401) if the signature is invalid. |
| 84 | + Raise HTTPException (403) if the signature is not set. |
| 85 | + """ |
| 86 | + |
| 87 | + signature = request.headers.get(self._header_name) |
| 88 | + if not signature: |
| 89 | + raise HTTPException(status_code=403, detail="Signature not set") |
| 90 | + if not self._compare_signature_fn( |
| 91 | + self._webhook_key, await request.body(), signature |
| 92 | + ): |
| 93 | + raise HTTPException(status_code=401, detail="Invalid signature") |
0 commit comments