poc2_worked

tests/test_notifier.py

@@ -0,0 +1,46 @@
+import json
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+from app.notifier import WebhookNotifier
+
+
+class WebhookHandler(BaseHTTPRequestHandler):
+    received = []
+
+    def do_POST(self):
+        length = int(self.headers.get("Content-Length", "0"))
+        WebhookHandler.received.append(json.loads(self.rfile.read(length)))
+        self.send_response(204)
+        self.send_header("Content-Length", "0")
+        self.end_headers()
+
+    def log_message(self, fmt, *args):
+        return
+
+
+def test_high_risk_webhook_is_async_and_rate_limited():
+    WebhookHandler.received = []
+    server = ThreadingHTTPServer(("127.0.0.1", 0), WebhookHandler)
+    threading.Thread(target=server.serve_forever, daemon=True).start()
+    notifier = WebhookNotifier(f"http://127.0.0.1:{server.server_port}/hook", 80, 2)
+    notifier.start()
+    try:
+        low = {"id": 1, "risk_score": 60, "subject_ip": "192.168.88.10"}
+        high = {"id": 2, "risk_score": 85, "subject_ip": "192.168.88.20", "stages": ["recon", "lateral-movement"]}
+        evidence = {"kind": "alert", "stage": "lateral-movement", "risk": 85, "summary": "test"}
+        notifier.notify(low, evidence)
+        notifier.notify(high, evidence)
+        notifier.notify(high, evidence)
+        deadline = time.time() + 2
+        while len(WebhookHandler.received) < 1 and time.time() < deadline:
+            time.sleep(0.02)
+        assert len(WebhookHandler.received) == 1
+        assert WebhookHandler.received[0]["incident"]["risk_score"] == 85
+        while notifier.status()["sent"] < 1 and time.time() < deadline:
+            time.sleep(0.01)
+        assert notifier.status()["sent"] == 1
+    finally:
+        notifier.stop()
+        server.shutdown(); server.server_close()