import asyncio import hashlib import hmac import json import time import random from collections import deque from dataclasses import dataclass from typing import Dict, Optional import aiohttp from aiohttp import web # --- CONFIGURATION --- SHARED_SECRET = "super_secret_twitch_key" DISCORD_WEBHOOK_URL = "https://discord.com/api/webhooks/mock" RATE_LIMIT_CAPACITY = 5 RATE_LIMIT_REFILL_PER_SEC = 5 IDEMPOTENCY_TTL_SECONDS = 300 # 5 minutes # --- MODELS & STATE --- @dataclass class DiscordResponse: status: int retry_after: Optional[float] = None class TokenBucket: """Implements the Token Bucket algorithm for rate limiting.""" def __init__(self, capacity: float, refill_rate: float): self.capacity = capacity self.refill_rate = refill_rate self.tokens = capacity self.last_update = time.monotonic() self.lock = asyncio.Lock() async def consume(self): async with self.lock: while self.tokens < 1: # Refill tokens based on time passed now = time.monotonic() passed = now - self.last_update self.tokens = min(self.capacity, self.tokens + (passed * self.refill_rate)) self.last_update = now if self.tokens < 1: # Wait for enough time to pass to get at least 1 token wait_time = (1 - self.tokens) / self.refill_rate await asyncio.sleep(wait_time) self.tokens -= 1 class IdempotencyStore: """In-memory store to prevent duplicate processing.""" def __init__(self, ttl: int): self.store: Dict[str, float] = {} self.ttl = ttl def is_duplicate(self, event_id: str) -> bool: now = time.time() if event_id in self.store: if now - self.store[event_id] < self.ttl: return True self.store[event_id] = now return False async def cleanup_loop(self): """Background task to evict stale entries.""" while True: await asyncio.sleep(60) now = time.time() expired = [k for k, v in self.store.items() if now - v > self.ttl] for k in expired: del self.store[k] # --- MOCK DISCORD SERVICE --- async def discord_send(payload: dict) -> DiscordResponse: """ Mocks a Discord webhook. Succeeds ~95% of the time. Returns 429 occasionally. """ # Simulate network latency await asyncio.sleep(0.05) rand = random.random() if rand < 0.05: # 5% chance of rate limit return DiscordResponse(status=429, retry_after=1.0) elif rand < 0.07: # 2% chance of other error return DiscordResponse(status=500) # Success return DiscordResponse(status=204) # --- CORE BRIDGE LOGIC --- class WebhookBridge: def __init__(self): self.idempotency = IdempotencyStore(IDEMPOTENCY_TTL_SECONDS) self.limiter = TokenBucket(RATE_LIMIT_CAPACITY, RATE_LIMIT_REFILL_PER_SEC) self.queue = asyncio.Queue() self.forward_count = 0 # For testing assertions async def worker(self): """Background worker that processes the queue and respects rate limits.""" while True: payload = await self.queue.get() try: await self._process_forward(payload) except Exception as e: print(f"[Worker Error] {e}") finally: self.queue.task_done() async def _process_forward(self, payload: dict): """Handles the actual POST to Discord with retry logic.""" await self.limiter.consume() resp = await discord_send(payload) if resp.status == 429: print(f"[Rate Limit] Discord returned 429. Retrying after {resp.retry_after}s") await asyncio.sleep(resp.retry_after or 1) # Retry once resp = await discord_send(payload) if resp.status in (200, 204): self.forward_count += 1 else: print(f"[Error] Discord failed with status {resp.status}") async def handle_webhook(self, request: web.Request) -> web.Response: """HTTP Handler for POST /webhook.""" # 1. Verify Signature signature_header = request.headers.get("X-Signature") if not signature_header: return web.Response(text="Missing X-Signature", status=401) body = await request.read() expected_sig = hmac.new( SHARED_SECRET.encode(), body, hashlib.sha256 ).hexdigest() if not hmac.compare_digest(expected_sig, signature_header): return web.Response(text="Invalid Signature", status=401) # 2. Parse JSON try: data = json.loads(body) event_id = data.get("event_id") if not event_id: return web.Response(text="Missing event_id", status=400) except json.JSONDecodeError: return web.Response(text="Invalid JSON", status=400) # 3. Idempotency Check if self.idempotency.is_duplicate(event_id): return web.Response(text="Duplicate event", status=200) # 4. Enqueue for forwarding await self.queue.put(data) return web.Response(text="Accepted", status=202) # --- TEST SUITE --- async def run_tests(bridge: WebhookBridge, server_url: str): print("\n--- Starting Tests ---") async with aiohttp.ClientSession() as session: # Helper to sign requests def sign(payload_dict): body = json.dumps(payload_dict).encode() sig = hmac.new(SHARED_SECRET.encode(), body, hashlib.sha256).hexdigest() return body, sig # a) Valid Request print("Test A: Valid signature...") payload_a = {"event_id": "evt_1", "type": "chat", "data": {"msg": "hi"}} body, sig = sign(payload_a) async with session.post(f"{server_url}/webhook", data=body, headers={"X-Signature": sig}) as r: assert r.status == 202 await asyncio.sleep(0.5) # Wait for worker assert bridge.forward_count == 1 print("✅ Passed") # b) Tampered Signature print("Test B: Tampered signature...") payload_b = {"event_id": "evt_2", "type": "chat"} body, _ = sign(payload_b) async with session.post(f"{server_url}/webhook", data=body, headers={"X-Signature": "wrong_sig"}) as r: assert r.status == 401 print("✅ Passed") # c) Idempotency (Replay) print("Test C: Replay event_id...") payload_c = {"event_id": "evt_1", "type": "chat"} # Same ID as Test A body, sig = sign(payload_c) async with session.post(f"{server_url}/webhook", data=body, headers={"X-Signature": sig}) as r: assert r.status == 200 # Returns 200 for duplicates per requirements await asyncio.sleep(0.5) assert bridge.forward_count == 1 # Count should NOT have increased print("✅ Passed") # d) Rate Limiting (Burst) print("Test D: Burst > 5 events/sec...") bridge.forward_count = 0 # Reset count for this test start_time = time.monotonic() # Send 10 events rapidly tasks = [] for i in range(10): p = {"event_id": f"burst_{i}", "type": "chat"} b, s = sign(p) tasks.append(session.post(f"{server_url}/webhook", data=b, headers={"X-Signature": s})) await asyncio.gather(*tasks) # Wait for all to be processed by worker await bridge.queue.join() duration = time.monotonic() - start_time # 10 events at 5/sec should take at least ~1.8-2.0 seconds # (First 5 instant, next 5 delayed by 1s each) print(f" Burst of 10 took {duration:.2f}s") assert duration >= 1.0 print("✅ Passed") print("--- All Tests Passed Successfully ---\n") # --- MAIN ENTRYPOINT --- async def main(): bridge = WebhookBridge() app = web.Application() app.router.add_post('/webhook', bridge.handle_webhook) runner = web.AppRunner(app) await runner.setup() site = web.TCPSite(runner, 'localhost', 8080) await site.start() print("Server started at http://localhost:8080") print("Press Ctrl+C to stop.") # Start background tasks worker_task = asyncio.create_task(bridge.worker()) cleanup_task = asyncio.create_task(bridge.idempotency.cleanup_loop()) # Run tests automatically try: await run_tests(bridge, "http://localhost:8080") except AssertionError as e: print(f"❌ Test Failed!") raise e except Exception as e: print(f"❌ Error during tests: {e}") raise e # Keep running until interrupted try: while True: await asyncio.sleep(3600) except asyncio.CancelledError: pass finally: worker_task.cancel() cleanup_task.cancel() await runner.cleanup() if __name__ == "__main__": try: asyncio.run(main()) except KeyboardInterrupt: pass