47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
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()
|