import json import os import sys import tempfile import threading import unittest from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer HERE = os.path.dirname(os.path.abspath(__file__)) ROOT = os.path.dirname(HERE) for path in (ROOT, HERE): if path not in sys.path: sys.path.insert(0, path) from graylog_ops import watch # noqa: E402 from graylog_ops.behavior import notify_fused # noqa: E402 from graylog_ops.notify import Notifier, load_notifier # noqa: E402 from test_watch import _StubClient # noqa: E402 class _CaptureNotifier: def __init__(self): self.batches = [] self.posts = [] def send(self, entries): self.batches.append(entries) def link_for(self, ip): return "" def post(self, text): self.posts.append(text) class _Receiver(BaseHTTPRequestHandler): received = [] def log_message(self, *args): pass def do_POST(self): length = int(self.headers.get("Content-Length") or 0) body = self.rfile.read(length).decode("utf-8") _Receiver.received.append(body) self.send_response(200) self.send_header("Content-Type", "application/json") self.end_headers() self.wfile.write(b'{"errcode":0}') class NotifierTests(unittest.TestCase): def _config(self, kind="generic", url="http://127.0.0.1:1/hook"): return {"type": kind, "url": url, "link_prefix": "http://example/panel/"} def test_payload_per_kind(self): text = "**新增恶意 IP**\n\n1.2.3.4" self.assertEqual( Notifier(self._config("wecom")).payload(text), {"msgtype": "markdown", "markdown": {"content": text}}) self.assertEqual( Notifier(self._config("dingtalk")).payload(text)["msgtype"], "markdown") self.assertEqual( Notifier(self._config("feishu")).payload(text)["msg_type"], "text") self.assertEqual(Notifier(self._config("slack")).payload(text), {"text": text}) generic = Notifier(self._config("generic")).payload(text) self.assertEqual(generic["text"], text) self.assertEqual(generic["title"], "新增恶意 IP") def test_render_includes_geo_rules_and_link(self): notifier = Notifier(self._config()) text = notifier.render([{ "ip": "203.0.113.5", "country_name": "美国", "region": "California", "score": 91.0, "total": 120, "first_seen": "2026-10-17T16:00:00Z", "rules_text": "Spring/中间件探测 4 次", }]) self.assertIn("203.0.113.5", text) self.assertIn("美国 · California", text) self.assertIn("Spring/中间件探测", text) self.assertIn("http://example/panel/#/ip/203.0.113.5", text) def test_send_to_http_endpoint(self): _Receiver.received = [] server = ThreadingHTTPServer(("127.0.0.1", 0), _Receiver) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() try: url = "http://127.0.0.1:%d/hook" % server.server_address[1] notifier = Notifier(self._config("generic", url)) notifier.send_test() self.assertEqual(len(_Receiver.received), 1) payload = json.loads(_Receiver.received[0]) self.assertIn("测试推送", payload["text"]) finally: server.shutdown() server.server_close() def test_dingtalk_signing(self): import base64 import hashlib import hmac import urllib.parse secret = "SECtest" url = "https://oapi.dingtalk.com/robot/send?access_token=abc" notifier = Notifier({"type": "dingtalk", "url": url, "secret": secret}) stamp = 1700000000000 expected = base64.b64encode(hmac.new( secret.encode(), ("%d\n%s" % (stamp, secret)).encode(), hashlib.sha256).digest()).decode() signed = notifier.signed_url(stamp) self.assertIn("timestamp=1700000000000", signed) self.assertIn("sign=" + urllib.parse.quote_plus(expected), signed) self.assertTrue(signed.startswith(url + "&")) def test_dingtalk_without_secret_keeps_url(self): url = "https://oapi.dingtalk.com/robot/send?access_token=abc" self.assertEqual(Notifier({"type": "dingtalk", "url": url}).signed_url(1), url) def test_dingtalk_signature(self): import base64 import hashlib import hmac import urllib.parse notifier = Notifier({ "type": "dingtalk", "url": "https://oapi.dingtalk.com/robot/send?access_token=abc", "secret": "SECtest-secret", }) url = notifier.signed_url() self.assertIn("timestamp=", url) self.assertIn("sign=", url) query = urllib.parse.parse_qs(urllib.parse.urlparse(url).query) timestamp = query["timestamp"][0] expected = base64.b64encode(hmac.new( b"SECtest-secret", ("%s\nSECtest-secret" % timestamp).encode("utf-8"), hashlib.sha256).digest()).decode("ascii") self.assertEqual(query["sign"][0], expected) unsigned = Notifier({"type": "dingtalk", "url": "https://example/hook?access_token=abc"}) self.assertEqual(unsigned.signed_url(), "https://example/hook?access_token=abc") def test_load_notifier(self): with tempfile.TemporaryDirectory() as directory: path = os.path.join(directory, "notify.json") json.dump(self._config("wecom"), open(path, "w", encoding="utf-8")) notifier = load_notifier(path) self.assertEqual(notifier.kind, "wecom") class UpsertStatusTests(unittest.TestCase): def test_upsert_returns_new_upgraded_unchanged(self): with tempfile.TemporaryDirectory() as directory: store = watch.Store(os.path.join(directory, "ips.db")) try: base = {"ip": "203.0.113.9", "total": 10, "rules": {}, "reasons": []} self.assertEqual(store.upsert(dict(base, tier=watch.TIER_OBSERVE, score=5.0), "2026-10-17T16:00:00Z"), "new") self.assertEqual(store.upsert(dict(base, tier=watch.TIER_OBSERVE, score=5.0), "2026-10-17T16:01:00Z"), "unchanged") self.assertEqual(store.upsert(dict(base, tier=watch.TIER_MALICIOUS, score=80.0), "2026-10-17T16:02:00Z"), "upgraded") self.assertEqual(store.upsert(dict(base, tier=watch.TIER_MALICIOUS, score=80.0), "2026-10-17T16:03:00Z"), "unchanged") finally: store.close() class DetectorNotificationTests(unittest.TestCase): def test_new_malicious_notifies_once(self): with tempfile.TemporaryDirectory() as directory: store = watch.Store(os.path.join(directory, "ips.db")) try: probe = watch.Detector(None, store, dict(watch.DEFAULT_FIELDS)) queries = {rule["name"]: rule["query"] for rule in probe.rules} by_query = {"*": [("203.0.113.5", 40)], queries["probe_java"]: [("203.0.113.5", 40)]} for name, query in queries.items(): by_query.setdefault(query, []) capture = _CaptureNotifier() detector = watch.Detector(_StubClient(by_query), store, dict(watch.DEFAULT_FIELDS), window=300) detector.run_once() store.sync_fused() # first sync suppresses historical rows; make it a fresh detection with store.lock, store.conn: store.conn.execute( "UPDATE ip_fused SET last_pushed_at=NULL") notify_fused(store, capture) self.assertEqual(len(capture.posts), 1) self.assertIn("203.0.113.5", capture.posts[0]) detector.run_once() store.sync_fused() notify_fused(store, capture) self.assertEqual(len(capture.posts), 1) # 24h cooldown finally: store.close() if __name__ == "__main__": unittest.main()