Add utility functions for webhooks

Author: rolodato

flagsmith/__init__.py

@@ -1,3 +1,4 @@
+from . import webhooks
 from .flagsmith import Flagsmith
 
-__all__ = ("Flagsmith",)
+__all__ = ("Flagsmith", "webhooks")

flagsmith/webhooks.py

@@ -0,0 +1,40 @@
+import hashlib
+import hmac
+
+
+def generate_signature(
+    request_body: str | bytes,
+    shared_secret: str,
+) -> str:
+    """Generates a signature for a webhook request body using HMAC-SHA256.
+
+    :param request_body: The raw request body, as string or bytes.
+    :param shared_secret: The shared secret configured for this specific webhook.
+    :return: The hex-encoded signature.
+    """
+    if isinstance(request_body, str):
+        request_body = request_body.encode()
+
+    shared_secret_bytes = shared_secret.encode()
+
+    return hmac.new(
+        key=shared_secret_bytes,
+        msg=request_body,
+        digestmod=hashlib.sha256,
+    ).hexdigest()
+
+
+def verify_signature(
+    request_body: str | bytes,
+    received_signature: str,
+    shared_secret: str,
+) -> bool:
+    """Verifies a webhook's signature to determine if the request was sent by Flagsmith.
+
+    :param request_body: The raw request body, as string or bytes.
+    :param received_signature: The signature as received in the X-Flagsmith-Signature request header.
+    :param shared_secret: The shared secret configured for this specific webhook.
+    :return: True if the signature is valid, False otherwise.
+    """
+    expected_signature = generate_signature(request_body, shared_secret)
+    return hmac.compare_digest(expected_signature, received_signature)

tests/test_webhooks.py

@@ -0,0 +1,50 @@
+import json
+
+from flagsmith.webhooks import generate_signature, verify_signature
+
+
+def test_generate_signature():
+    # Given
+    request_body = json.dumps({"data": {"foo": 123}})
+    shared_secret = "shh"
+
+    # When
+    signature = generate_signature(request_body, shared_secret)
+
+    # Then
+    assert isinstance(signature, str)
+    assert len(signature) == 64  # SHA-256 hex digest is 64 characters
+
+
+def test_verify_signature_valid():
+    # Given
+    request_body = json.dumps({"data": {"foo": 123}})
+    shared_secret = "shh"
+
+    # When
+    signature = generate_signature(request_body, shared_secret)
+
+    # Then
+    assert verify_signature(
+        request_body=request_body,
+        received_signature=signature,
+        shared_secret=shared_secret,
+    )
+    # Test with bytes instead of str
+    assert verify_signature(
+        request_body=request_body.encode(),
+        received_signature=signature,
+        shared_secret=shared_secret,
+    )
+
+
+def test_verify_signature_invalid():
+    # Given
+    request_body = json.dumps({"event": "flag_updated", "data": {"id": 123}})
+
+    # Then
+    assert not verify_signature(
+        request_body=request_body,
+        received_signature="bad",
+        shared_secret="?",
+    )