"""Tests for the judgement layer (ban candidate scoring).""" import io import json import os import sys import tempfile import unittest from contextlib import redirect_stderr, redirect_stdout HERE = os.path.dirname(os.path.abspath(__file__)) ROOT = os.path.dirname(HERE) sys.path.insert(0, ROOT) sys.path.insert(0, HERE) from fake_graylog import JUDGE_MESSAGES, FakeGraylog # noqa: E402 from graylog_ops import cli # noqa: E402 class JudgeTests(unittest.TestCase): @classmethod def setUpClass(cls): cls.server = FakeGraylog(messages=JUDGE_MESSAGES).start() @classmethod def tearDownClass(cls): cls.server.stop() def run_cli(self, *argv, json_mode=True): args = ["--url", self.server.url, "--token", "testtoken"] if json_mode: args.append("--json") args.extend(argv) stdout, stderr = io.StringIO(), io.StringIO() with redirect_stdout(stdout), redirect_stderr(stderr): code = cli.main(args) return code, stdout.getvalue(), stderr.getvalue() def judged(self, *argv): code, out, err = self.run_cli("judge", "--range", "300", *argv) self.assertEqual(code, 0, err) return json.loads(out) # ------------------------------------------------------------ behaviour def test_bruteforce_client_is_recommended_for_ban(self): data = self.judged() rows = {row["ip"]: row for row in data["candidates"]} self.assertIn("203.0.113.66", rows) row = rows["203.0.113.66"] self.assertEqual(row["action"], "ban_1h") self.assertGreaterEqual(row["score"], 75) self.assertIn("auth_fail", row["signals"]) self.assertTrue(any("失败请求" in reason for reason in row["reasons"])) def test_scanner_client_is_recommended_for_ban(self): data = self.judged() row = {item["ip"]: item for item in data["candidates"]}["192.0.2.99"] self.assertEqual(row["action"], "ban_1h") self.assertIn("scanner_ua", row["signals"]) self.assertIn("sensitive_path", row["signals"]) def test_busy_but_legitimate_client_is_not_banned(self): data = self.judged("--all") row = {item["ip"]: item for item in data["candidates"]}["198.51.100.200"] self.assertIn(row["action"], ("ignore", "observe")) self.assertFalse(row["evidence"]) def test_internal_addresses_are_excluded_by_default(self): data = self.judged() reasons = {row["ip"]: row["reason"] for row in data["excluded"]} self.assertIn("10.180.5.5", reasons) self.assertIn("内部", reasons["10.180.5.5"]) self.assertNotIn("10.180.5.5", {row["ip"] for row in data["candidates"]}) def test_include_internal_flag(self): data = self.judged("--include-internal", "--all") ips = {row["ip"] for row in data["candidates"]} self.assertIn("10.180.5.5", ips) def test_allowlist_file(self): with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as handle: handle.write("# 深信服 AD / VIP\n10.180.254.17\n203.0.113.66\n") path = handle.name try: data = self.judged("--allow", path) excluded = {row["ip"]: row["reason"] for row in data["excluded"]} self.assertIn("203.0.113.66", excluded) self.assertIn("白名单", excluded["203.0.113.66"]) self.assertNotIn("203.0.113.66", {row["ip"] for row in data["candidates"]}) finally: os.unlink(path) def test_explain_attaches_samples(self): data = self.judged("--explain", "2") sampled = [row for row in data["candidates"] if row.get("samples")] self.assertTrue(sampled) self.assertLessEqual(len(sampled[0]["samples"]), 3) def test_verify_downgrades_automation_client(self): """A steady single-endpoint poller with 403s must not be auto-banned.""" automation = [{ "timestamp": "2026-10-17T15:%02d:%02d.000Z" % (index // 60, index % 60), "source": "10.180.254.17", "client_ip": "203.0.113.77", "status": 403, "request_method": "POST", "request_uri": "/prod-api/biz/liveCourseChapter/msgSubscription", "http_user_agent": "okhttp/4.9", "message": '203.0.113.77 - "POST /prod-api/biz/x" 403', } for index in range(40)] server = FakeGraylog(messages=list(JUDGE_MESSAGES) + automation).start() try: stdout, stderr = io.StringIO(), io.StringIO() with redirect_stdout(stdout), redirect_stderr(stderr): code = cli.main(["--url", server.url, "--token", "testtoken", "--json", "judge", "--range", "300", "--verify", "10", "--all"]) self.assertEqual(code, 0, stderr.getvalue()) data = json.loads(stdout.getvalue()) row = {item["ip"]: item for item in data["candidates"]}["203.0.113.77"] self.assertEqual(row["action"], "observe") self.assertTrue(row["verify"]["automation_suspected"]) self.assertFalse(row["verify"]["login_like"]) finally: server.stop() def test_verify_keeps_login_bruteforce(self): """Failures against a login path stay bannable even if concentrated.""" data = self.judged("--verify", "5") row = {item["ip"]: item for item in data["candidates"]}["203.0.113.66"] self.assertEqual(row["action"], "ban_1h") self.assertTrue(row["verify"]["login_like"]) def test_profile_roundtrip(self): with tempfile.TemporaryDirectory() as directory: profile = os.path.join(directory, "profile.json") code, _, err = self.run_cli("discover", "--range", "3600", "--save-profile", profile) self.assertEqual(code, 0, err) with open(profile, "r", encoding="utf-8") as handle: saved = json.load(handle) self.assertEqual(saved["fields"]["client_ip"], "client_ip") self.assertEqual(saved["fields"]["status"], "status") self.assertEqual(saved["fields"]["uri"], "request_uri") data = self.judged("--profile", profile) self.assertTrue(data["candidates"]) self.assertEqual(data["field"], "client_ip") # -------------------------------------------------------------- output def test_human_output(self): code, out, err = self.run_cli("judge", json_mode=False) self.assertEqual(code, 0, err) self.assertIn("信号命中", out) self.assertIn("本命令只读", out) self.assertIn("ban_1h", out) def test_signal_queries_are_reported(self): data = self.judged() names = {stat["name"]: stat for stat in data["signal_stats"]} self.assertIn("auth_fail", names) self.assertGreater(names["auth_fail"]["ips_matched"], 0) self.assertGreater(names["scanner_ua"]["ips_matched"], 0) self.assertIn("query", names["auth_fail"]) if __name__ == "__main__": unittest.main()