大活

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

@@ -0,0 +1,203 @@
+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")