EPEEAIKit/art-agent/backend/venv/Lib/site-packages/replicate/webhook.py

import base64
import hmac
from hashlib import sha256
from typing import (
    TYPE_CHECKING,
    Dict,
    Optional,
    overload,
)

from replicate.resource import Namespace, Resource

if TYPE_CHECKING:
    import httpx


class WebhookSigningSecret(Resource):
    """
    A webhook signing secret.
    """

    key: str


class WebhookValidationError(ValueError):
    """Base exception for webhook validation errors."""


class MissingWebhookHeaderError(WebhookValidationError):
    """Exception raised when a required webhook header is missing."""


class InvalidSecretKeyError(WebhookValidationError):
    """Exception raised when the secret key format is invalid."""


class MissingWebhookBodyError(WebhookValidationError):
    """Exception raised when the webhook body is missing."""


class InvalidTimestampError(WebhookValidationError):
    """Exception raised when the webhook timestamp is invalid or outside the tolerance."""


class InvalidSignatureError(WebhookValidationError):
    """Exception raised when the webhook signature is invalid."""


class Webhooks(Namespace):
    """
    Namespace for operations related to webhooks.
    """

    @property
    def default(self) -> "Webhooks.Default":
        """
        Namespace for operations related to the default webhook.
        """

        return self.Default(self._client)

    class Default(Namespace):
        """
        Namespace for operations related to the default webhook.
        """

        def secret(self) -> WebhookSigningSecret:
            """
            Get the default webhook signing secret.

            Returns:
                WebhookSigningSecret: The default webhook signing secret.
            """

            resp = self._client._request("GET", "/v1/webhooks/default/secret")
            return WebhookSigningSecret(**resp.json())

        async def async_secret(self) -> WebhookSigningSecret:
            """
            Get the default webhook signing secret.

            Returns:
                WebhookSigningSecret: The default webhook signing secret.
            """

            resp = await self._client._async_request(
                "GET", "/v1/webhooks/default/secret"
            )
            return WebhookSigningSecret(**resp.json())

    @overload
    @staticmethod
    def validate(
        request: "httpx.Request",
        secret: WebhookSigningSecret,
        tolerance: Optional[int] = None,
    ) -> bool: ...

    @overload
    @staticmethod
    def validate(
        headers: Dict[str, str],
        body: str,
        secret: WebhookSigningSecret,
        tolerance: Optional[int] = None,
    ) -> bool: ...

    @staticmethod
    def validate(  # type: ignore # pylint: disable=too-many-branches,too-many-locals
        request: Optional["httpx.Request"] = None,
        headers: Optional[Dict[str, str]] = None,
        body: Optional[str] = None,
        secret: Optional[WebhookSigningSecret] = None,
        tolerance: Optional[int] = None,
    ) -> None:
        """
        Validate the signature from an incoming webhook request using the provided secret.

        Args:
            request (httpx.Request): The request object.
            headers (Dict[str, str]): The request headers.
            body (str): The request body.
            secret (WebhookSigningSecret): The webhook signing secret.
            tolerance (Optional[int]): Maximum allowed time difference (in seconds) between the current time and the webhook timestamp.

        Returns:
            None: If the request is valid.

        Raises:
            MissingWebhookHeaderError: If required webhook headers are missing.
            InvalidSecretKeyError: If the secret key format is invalid.
            MissingWebhookBodyError: If the webhook body is missing.
            InvalidTimestampError: If the webhook timestamp is invalid or outside the tolerance.
            InvalidSignatureError: If the webhook signature is invalid.
        """

        if not secret:
            raise ValueError("Missing webhook signing secret")

        if request and any([headers, body]):
            raise ValueError("Only one of request or headers/body can be provided")

        if request and request.headers:
            webhook_id = request.headers.get("webhook-id")
            timestamp = request.headers.get("webhook-timestamp")
            signature = request.headers.get("webhook-signature")
            body = request.content.decode("utf-8")
        else:
            if not headers:
                raise MissingWebhookHeaderError("Missing webhook headers")

            # Convert headers to case-insensitive dictionary
            headers = {k.lower(): v for k, v in headers.items()}

            webhook_id = headers.get("webhook-id")
            timestamp = headers.get("webhook-timestamp")
            signature = headers.get("webhook-signature")

        if not webhook_id:
            raise MissingWebhookHeaderError("Missing webhook id")
        if not timestamp:
            raise MissingWebhookHeaderError("Missing webhook timestamp")
        if not signature:
            raise MissingWebhookHeaderError("Missing webhook signature")
        if not body:
            raise MissingWebhookBodyError("Missing webhook body")

        if tolerance is not None:
            import time  # pylint: disable=import-outside-toplevel

            current_time = int(time.time())
            webhook_time = int(timestamp)
            time_difference = abs(current_time - webhook_time)
            if time_difference > tolerance:
                raise InvalidTimestampError(
                    f"Webhook timestamp is outside the allowed tolerance of {tolerance} seconds"
                )

        signed_content = f"{webhook_id}.{timestamp}.{body}"

        key_parts = secret.key.split("_")
        if len(key_parts) != 2:
            raise InvalidSecretKeyError(f"Invalid secret key format: {secret.key}")

        secret_bytes = base64.b64decode(key_parts[1])

        h = hmac.new(secret_bytes, signed_content.encode(), sha256)
        computed_signature = h.digest()

        valid = False
        for sig in signature.split():
            sig_parts = sig.split(",")
            if len(sig_parts) < 2:
                raise InvalidSignatureError(f"Invalid signature format: {sig}")

            sig_bytes = base64.b64decode(sig_parts[1])

            if hmac.compare_digest(sig_bytes, computed_signature):
                valid = True
                break

        if not valid:
            raise InvalidSignatureError("Webhook signature is invalid")