522 lines
25 KiB
Python
522 lines
25 KiB
Python
"""Tests for the continuous detector, the SQLite store and the dashboard."""
|
|
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
import urllib.error
|
|
import urllib.request
|
|
from http.server import ThreadingHTTPServer
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
ROOT = os.path.dirname(HERE)
|
|
sys.path.insert(0, ROOT)
|
|
|
|
from graylog_ops import watch # noqa: E402
|
|
|
|
|
|
class ClassifyTests(unittest.TestCase):
|
|
def test_sensitive_probe_is_malicious(self):
|
|
tier, score, reasons = watch.classify(44, {"probe_java": 44})
|
|
self.assertEqual(tier, watch.TIER_MALICIOUS)
|
|
self.assertGreaterEqual(score, 60)
|
|
self.assertTrue(any("探测" in reason for reason in reasons))
|
|
|
|
def test_single_sensitive_hit_is_enough(self):
|
|
tier, score, _ = watch.classify(18, {"probe_java": 18})
|
|
self.assertEqual(tier, watch.TIER_MALICIOUS)
|
|
|
|
def test_dotbot_crawler_is_malicious(self):
|
|
tier, score, reasons = watch.classify(352, {"bot_ua": 352, "not_found": 176})
|
|
self.assertEqual(tier, watch.TIER_MALICIOUS)
|
|
self.assertTrue(any("404" in reason for reason in reasons))
|
|
|
|
def test_login_burst_alone_is_only_observed(self):
|
|
"""Students logging in repeatedly must not look suspicious."""
|
|
tier, score, _ = watch.classify(2253, {"login_burst": 660})
|
|
self.assertEqual(tier, watch.TIER_OBSERVE)
|
|
|
|
def test_login_attack_plus_burst_is_malicious(self):
|
|
tier, score, _ = watch.classify(666, {"login_attack": 333, "login_burst": 666})
|
|
self.assertEqual(tier, watch.TIER_MALICIOUS)
|
|
|
|
def test_busy_clean_ip_scores_zero(self):
|
|
tier, score, reasons = watch.classify(5000, {})
|
|
self.assertEqual(score, 0.0)
|
|
self.assertEqual(reasons, [])
|
|
|
|
def test_404_ratio_needs_volume(self):
|
|
_, score, _ = watch.classify(10, {"not_found": 5})
|
|
self.assertEqual(score, 0.0)
|
|
|
|
|
|
class _StubClient:
|
|
"""Returns canned counts keyed by the exact query string."""
|
|
|
|
def __init__(self, by_query, totals=None):
|
|
self.by_query = by_query
|
|
self.totals = totals or {}
|
|
|
|
def terms(self, field, query, size, range_seconds, filter_str=None):
|
|
if query.startswith("client_ip:("):
|
|
ips = [part.strip().strip('"')
|
|
for part in query[len("client_ip:("):-1].split(" OR ")]
|
|
items = [(ip, self.totals[ip]) for ip in ips if ip in self.totals]
|
|
else:
|
|
items = self.by_query.get(query, [])
|
|
return {"terms": [{"term": ip, "count": count} for ip, count in items],
|
|
"total": sum(count for _, count in items)}
|
|
|
|
def search_views_messages(self, query="*", range_seconds=300, limit=5, offset=0,
|
|
filter_str=None, from_ts=None, to_ts=None,
|
|
timeout_ms=60000, sort_order="DESC"):
|
|
return []
|
|
|
|
|
|
class StoreTests(unittest.TestCase):
|
|
def test_upsert_list_forget(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
try:
|
|
record = {"ip": "203.0.113.9", "tier": watch.TIER_MALICIOUS,
|
|
"score": 88.0, "total": 120,
|
|
"rules": {"sensitive_probe": 4}, "reasons": ["敏感路径探测 4 次"]}
|
|
store.upsert(record, "2026-10-17T16:00:00Z")
|
|
store.upsert(dict(record, score=90.0), "2026-10-17T16:01:00Z")
|
|
rows = store.list()
|
|
self.assertEqual(len(rows), 1)
|
|
self.assertEqual(rows[0]["score"], 90.0)
|
|
self.assertEqual(rows[0]["hits"], 2)
|
|
self.assertEqual(rows[0]["rules"]["sensitive_probe"], 4)
|
|
|
|
self.assertEqual(store.stats()["by_tier"][watch.TIER_MALICIOUS], 1)
|
|
self.assertEqual(store.stats()["events"], 1)
|
|
|
|
store.upsert({"ip": "198.51.100.7", "tier": watch.TIER_SUSPICIOUS,
|
|
"score": 40.0, "total": 10, "rules": {}, "reasons": []},
|
|
"2026-10-17T16:02:00Z")
|
|
self.assertEqual(len(store.list(tier=watch.TIER_SUSPICIOUS)), 1)
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
class PermanentRetentionTests(unittest.TestCase):
|
|
def _store(self, directory):
|
|
return watch.Store(os.path.join(directory, "ips.db"))
|
|
|
|
def test_ever_malicious_survives_downgrade(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = self._store(directory)
|
|
try:
|
|
base = {"ip": "203.0.113.9", "total": 10, "rules": {}, "reasons": []}
|
|
store.upsert(dict(base, tier=watch.TIER_MALICIOUS, score=80.0),
|
|
"2026-10-17T16:00:00Z")
|
|
store.upsert(dict(base, tier=watch.TIER_OBSERVE, score=5.0),
|
|
"2026-10-17T16:01:00Z")
|
|
row = store.get("203.0.113.9")
|
|
self.assertEqual(row["tier"], watch.TIER_OBSERVE)
|
|
self.assertEqual(row["ever_malicious"], 1)
|
|
self.assertEqual(row["last_malicious_at"], "2026-10-17T16:00:00Z")
|
|
core = store.list(scope="core")
|
|
self.assertEqual([item["ip"] for item in core], ["203.0.113.9"])
|
|
finally:
|
|
store.close()
|
|
|
|
def test_forget_keeps_convicted_ips(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = self._store(directory)
|
|
try:
|
|
store.upsert({"ip": "203.0.113.9", "tier": watch.TIER_MALICIOUS,
|
|
"score": 80.0, "total": 10, "rules": {},
|
|
"reasons": []}, "2026-10-01T00:00:00Z")
|
|
store.upsert({"ip": "198.51.100.7", "tier": watch.TIER_OBSERVE,
|
|
"score": 5.0, "total": 10, "rules": {},
|
|
"reasons": []}, "2026-10-01T00:00:00Z")
|
|
store.forget(3600, time.time())
|
|
self.assertIsNotNone(store.get("203.0.113.9"))
|
|
self.assertIsNone(store.get("198.51.100.7"))
|
|
finally:
|
|
store.close()
|
|
|
|
def test_sorting_modes(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = self._store(directory)
|
|
try:
|
|
store.upsert({"ip": "203.0.113.2", "tier": watch.TIER_MALICIOUS,
|
|
"score": 70.0, "total": 10, "rules": {},
|
|
"reasons": []}, "2026-10-17T16:00:00Z")
|
|
store.upsert({"ip": "203.0.113.1", "tier": watch.TIER_MALICIOUS,
|
|
"score": 90.0, "total": 5, "rules": {},
|
|
"reasons": []}, "2026-10-17T17:00:00Z")
|
|
first = [row["ip"] for row in store.list(sort="first_seen", order="asc")]
|
|
self.assertEqual(first, ["203.0.113.2", "203.0.113.1"])
|
|
recent = [row["ip"] for row in store.list(sort="last_malicious",
|
|
order="desc")]
|
|
self.assertEqual(recent, ["203.0.113.1", "203.0.113.2"])
|
|
score = [row["ip"] for row in store.list(sort="score", order="desc")]
|
|
self.assertEqual(score, ["203.0.113.1", "203.0.113.2"])
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
class TextSanitisingTests(unittest.TestCase):
|
|
def test_strips_control_and_replacement_chars(self):
|
|
self.assertEqual(watch._clean_text("a\x00b\ufffdc"), "a·b·c")
|
|
|
|
def test_scrubs_binary_header_field(self):
|
|
self.assertEqual(watch._scrub_binary("\ufffd\ufffdk\ufffdU\ufffd\ufffd"),
|
|
"[非 UTF-8 数据]")
|
|
self.assertEqual(watch._scrub_binary("GET"), "GET")
|
|
|
|
|
|
class DetectorTests(unittest.TestCase):
|
|
def _detector(self, by_query, store):
|
|
fields = dict(watch.DEFAULT_FIELDS)
|
|
probe = watch.Detector(None, store, fields)
|
|
queries = {rule["name"]: rule["query"] for rule in probe.rules}
|
|
mapping = dict(by_query)
|
|
mapping.setdefault(queries["not_found"], [])
|
|
return watch.Detector(_StubClient(mapping), store, fields, window=300)
|
|
|
|
def test_automation_is_downgraded(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))
|
|
q = {rule["name"]: rule["query"] for rule in probe.rules}
|
|
by_query = {
|
|
"*": [("124.70.2.111", 666)],
|
|
q["login_attack"]: [("124.70.2.111", 333)],
|
|
q["login_burst"]: [("124.70.2.111", 666)],
|
|
q["not_found"]: [],
|
|
q["probe_java"]: [],
|
|
q["bad_ua"]: [],
|
|
q["bot_ua"]: [],
|
|
q["probe_config"]: [],
|
|
q["probe_cms"]: [],
|
|
q["probe_iot"]: [],
|
|
q["probe_devops"]: [],
|
|
q["probe_webshell"]: [],
|
|
q["exploit_payload"]: [],
|
|
'client_ip:"124.70.2.111"': [
|
|
("/prod-api/biz/liveCourseChapter/msgSubscription", 666)],
|
|
}
|
|
detector = watch.Detector(_StubClient(by_query), store,
|
|
dict(watch.DEFAULT_FIELDS), window=300)
|
|
result = detector.run_once()
|
|
self.assertEqual(result["automation"], 1)
|
|
row = store.get("124.70.2.111")
|
|
self.assertEqual(row["tier"], watch.TIER_AUTOMATION)
|
|
self.assertTrue(any("单一接口" in r for r in row["reasons"]))
|
|
finally:
|
|
store.close()
|
|
|
|
def test_low_volume_scanner_is_still_caught(self):
|
|
"""A scanner with 18 requests must not need a top-volume slot."""
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
try:
|
|
probe = watch.Detector(None, store, dict(watch.DEFAULT_FIELDS))
|
|
q = {rule["name"]: rule["query"] for rule in probe.rules}
|
|
by_query = {
|
|
"*": [("10.185.9.41", 5000), ("202.201.16.87", 4000)],
|
|
q["probe_java"]: [("5.188.86.234", 18)],
|
|
}
|
|
detector = watch.Detector(
|
|
_StubClient(by_query, totals={"5.188.86.234": 18}), store,
|
|
dict(watch.DEFAULT_FIELDS), window=300)
|
|
result = detector.run_once()
|
|
self.assertEqual(result["malicious"], 1)
|
|
self.assertEqual(store.get("5.188.86.234")["tier"], watch.TIER_MALICIOUS)
|
|
finally:
|
|
store.close()
|
|
|
|
def test_internal_and_allowlisted_are_skipped(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))
|
|
q = {rule["name"]: rule["query"] for rule in probe.rules}
|
|
by_query = {
|
|
"*": [("10.185.9.41", 5000), ("203.0.113.5", 30), ("198.51.100.9", 30)],
|
|
q["probe_java"]: [("203.0.113.5", 30), ("198.51.100.9", 30)],
|
|
}
|
|
detector = watch.Detector(
|
|
_StubClient(by_query), store, dict(watch.DEFAULT_FIELDS),
|
|
window=300, allowlist=lambda ip: ip == "198.51.100.9")
|
|
detector.run_once()
|
|
self.assertIsNone(store.get("10.185.9.41"))
|
|
self.assertIsNone(store.get("198.51.100.9"))
|
|
self.assertEqual(store.get("203.0.113.5")["tier"], watch.TIER_MALICIOUS)
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
class SampleTests(unittest.TestCase):
|
|
def test_sample_keeps_all_source_fields(self):
|
|
message = {
|
|
"timestamp": "2026-10-17T05:08:58.046Z",
|
|
"client_ip": "51.158.145.102", "client_port": 46990,
|
|
"server_ip": "10.185.9.90", "server_port": 80,
|
|
"host": "61.178.20.46:443", "bigip": "www.gsedu.edu.cn_2_4600",
|
|
"vs": "/Common/SmartEDU_https", "method": "GET",
|
|
"uri": "/vendor/phpunit/phpunit/src/Util/PHP/eval-stdin.php",
|
|
"status": "403", "ua": "libredtail-http",
|
|
"client_ip_country_code": "FR", "rt": "0.031",
|
|
"message": 'event_type="req",client_ip="51.158.145.102"',
|
|
}
|
|
sample = watch._sample(message)
|
|
for key in ("client_port", "server_port", "bigip", "vs", "rt",
|
|
"client_ip_country_code", "message"):
|
|
self.assertIn(key, sample)
|
|
self.assertEqual(sample["client_port"], "46990")
|
|
|
|
def test_sample_drops_empty_and_na_values(self):
|
|
sample = watch._sample({"client_ip": "1.2.3.4",
|
|
"client_ip_city_name": "N/A",
|
|
"referer": ""})
|
|
self.assertEqual(sample, {"client_ip": "1.2.3.4"})
|
|
|
|
|
|
class SnapshotTests(unittest.TestCase):
|
|
def test_compact_evidence_is_bounded(self):
|
|
evidence = {
|
|
"window": 89600,
|
|
"uris": [{"term": "/probe/%d" % i, "count": i} for i in range(30)],
|
|
"statuses": [{"term": str(code), "count": 1} for code in range(10)],
|
|
"methods": [{"term": "GET", "count": 1}],
|
|
"hosts": [{"term": "h", "count": 1}],
|
|
"server_ips": [{"term": "10.0.0.1", "count": 1}],
|
|
"vips": [{"term": "vip", "count": 1}],
|
|
"uas": [{"term": "ua", "count": 1}],
|
|
"samples": [{} for _ in range(5)],
|
|
"timeline": [{"timestamp": "t%d" % i, "count": 1} for i in range(200)],
|
|
"peak": {"timestamp": "t", "count": 3},
|
|
}
|
|
compact = watch._compact_evidence(evidence)
|
|
self.assertTrue(compact["compact"])
|
|
self.assertLessEqual(len(compact["uris"]), 10)
|
|
self.assertLessEqual(len(compact["timeline"]), 72)
|
|
self.assertEqual(compact["sample_count"], 5)
|
|
self.assertEqual(compact["window"], 89600)
|
|
|
|
|
|
class PathPrefixTests(unittest.TestCase):
|
|
def test_root_is_hidden_and_prefix_works(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
store.upsert({"ip": "203.0.113.9", "tier": watch.TIER_MALICIOUS,
|
|
"score": 91.0, "total": 120, "window": 3600,
|
|
"rules": {"probe_java": 4},
|
|
"reasons": ["探测 4 次"]}, "2026-10-17T16:00:00Z")
|
|
store.sync_fused()
|
|
server = ThreadingHTTPServer(
|
|
("127.0.0.1", 0),
|
|
watch.make_handler(store, 3600, 60, None, "/s/secret123"))
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
base = "http://127.0.0.1:%d" % server.server_address[1]
|
|
try:
|
|
with self.assertRaises(urllib.error.HTTPError) as ctx:
|
|
urllib.request.urlopen(base + "/", timeout=5)
|
|
self.assertEqual(ctx.exception.code, 404)
|
|
|
|
with urllib.request.urlopen(base + "/s/secret123/", timeout=5) as response:
|
|
page = response.read().decode("utf-8")
|
|
self.assertIn('const BASE = "/s/secret123"', page)
|
|
|
|
with urllib.request.urlopen(base + "/s/secret123/api/ips", timeout=5) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
self.assertEqual(payload["count"], 1)
|
|
|
|
with self.assertRaises(urllib.error.HTTPError):
|
|
urllib.request.urlopen(base + "/api/ips", timeout=5)
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
store.close()
|
|
|
|
|
|
class MetricsTests(unittest.TestCase):
|
|
def test_metrics_and_monitor(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
|
store.set_meta("detect_last_success", now)
|
|
store.set_meta("detect_last_error", "")
|
|
store.incr_meta("detect_rounds_total", 3)
|
|
store.incr_meta("detect_errors_total", 1)
|
|
store.set_meta("behavior_last_success", now)
|
|
store.set_meta("behavior_last_duration", 12.5)
|
|
store.set_meta("behavior_last_count", 7)
|
|
store.incr_meta("behavior_runs_total", 2)
|
|
store.incr_meta("push_sent_total", 5)
|
|
store.incr_meta("push_failed_total", 2)
|
|
store.upsert({"ip": "203.0.113.9", "tier": watch.TIER_MALICIOUS,
|
|
"score": 91.0, "total": 120, "window": 3600,
|
|
"rules": {"probe_java": 2},
|
|
"reasons": ["探测"]}, "2026-10-17T16:00:00Z")
|
|
store.sync_fused()
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0),
|
|
watch.make_handler(store, 3600, 60))
|
|
threading.Thread(target=server.serve_forever, daemon=True).start()
|
|
base = "http://127.0.0.1:%d" % server.server_address[1]
|
|
try:
|
|
with urllib.request.urlopen(base + "/metrics", timeout=5) as response:
|
|
text = response.read().decode("utf-8")
|
|
self.assertIn("ipsec_detect_rounds_total 3", text)
|
|
self.assertIn("ipsec_detect_errors_total 1", text)
|
|
self.assertIn("ipsec_behavior_last_duration_seconds 12.5", text)
|
|
self.assertIn("ipsec_fused_total 1", text)
|
|
self.assertIn("ipsec_push_sent_total 5", text)
|
|
self.assertIn("ipsec_graylog_up 1", text)
|
|
|
|
with urllib.request.urlopen(base + "/api/monitor", timeout=5) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
self.assertEqual(payload["graylog"], "ok")
|
|
self.assertEqual(payload["detect"]["rounds"], 3)
|
|
self.assertEqual(payload["behavior"]["last_duration"], 12.5)
|
|
self.assertEqual(payload["push"]["sent_total"], 5)
|
|
self.assertEqual(payload["list"]["fused_total"], 1)
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
store.close()
|
|
|
|
|
|
class PaginationTests(unittest.TestCase):
|
|
def test_pagination_and_search(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
for index, ip in enumerate(("203.0.113.9", "203.0.113.10",
|
|
"198.51.100.77")):
|
|
store.upsert({"ip": ip, "tier": watch.TIER_MALICIOUS,
|
|
"score": 90.0, "total": 100 + index,
|
|
"window": 3600, "first_seen":
|
|
"2026-10-17T16:0%d:00Z" % index,
|
|
"rules": {"probe_java": 2},
|
|
"reasons": ["敏感路径探测"]},
|
|
"2026-10-17T16:30:00Z")
|
|
store.sync_fused()
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0),
|
|
watch.make_handler(store, 3600, 60))
|
|
threading.Thread(target=server.serve_forever, daemon=True).start()
|
|
base = "http://127.0.0.1:%d" % server.server_address[1]
|
|
try:
|
|
with urllib.request.urlopen(
|
|
base + "/api/ips?page_size=2&page=2", timeout=5) as response:
|
|
page = json.loads(response.read().decode("utf-8"))
|
|
self.assertEqual(page["total"], 3)
|
|
self.assertEqual(page["pages"], 2)
|
|
self.assertEqual(page["page"], 2)
|
|
self.assertEqual(len(page["items"]), 1)
|
|
|
|
with urllib.request.urlopen(
|
|
base + "/api/ips?q=203.0.113", timeout=5) as response:
|
|
filtered = json.loads(response.read().decode("utf-8"))
|
|
self.assertEqual(filtered["total"], 2)
|
|
self.assertEqual(len(filtered["items"]), 2)
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
store.close()
|
|
|
|
|
|
class StatusIndicatorTests(unittest.TestCase):
|
|
def test_status_endpoint(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
|
|
class DetectorStub:
|
|
last_success = time.time() - 5
|
|
last_error = None
|
|
last_run = "2026-10-17T16:00:00Z"
|
|
|
|
server = ThreadingHTTPServer(
|
|
("127.0.0.1", 0),
|
|
watch.make_handler(store, 3600, 60, DetectorStub()))
|
|
threading.Thread(target=server.serve_forever, daemon=True).start()
|
|
base = "http://127.0.0.1:%d" % server.server_address[1]
|
|
try:
|
|
with urllib.request.urlopen(base + "/api/status", timeout=5) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
self.assertEqual(payload["graylog"], "ok")
|
|
|
|
DetectorStub.last_error = "connection refused"
|
|
DetectorStub.last_success = None
|
|
with urllib.request.urlopen(base + "/api/status", timeout=5) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
self.assertEqual(payload["graylog"], "error")
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
store.close()
|
|
|
|
|
|
class DashboardTests(unittest.TestCase):
|
|
def test_api_and_html(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
store.upsert({"ip": "203.0.113.9", "tier": watch.TIER_MALICIOUS,
|
|
"score": 91.0, "total": 120, "window": 3600,
|
|
"rules": {"sensitive_probe": 4},
|
|
"reasons": ["敏感路径探测 4 次"]}, "2026-10-17T16:00:00Z")
|
|
store.upsert({"ip": "198.51.100.77", "tier": watch.TIER_MALICIOUS,
|
|
"score": 70.0, "total": 653344, "window": 3600,
|
|
"risk": "shared_egress",
|
|
"rules": {"sensitive_probe": 4},
|
|
"reasons": ["疑似共享出口"]}, "2026-10-17T16:00:00Z")
|
|
store.sync_fused()
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0),
|
|
watch.make_handler(store, 3600, 60))
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
base = "http://127.0.0.1:%d" % server.server_address[1]
|
|
try:
|
|
with urllib.request.urlopen(base + "/", timeout=5) as response:
|
|
page = response.read().decode("utf-8")
|
|
self.assertIn("恶意 IP 清单", page)
|
|
self.assertIn("/api/ips", page)
|
|
|
|
with urllib.request.urlopen(base + "/api/ips", timeout=5) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
self.assertEqual(payload["count"], 2)
|
|
self.assertEqual({item["ip"] for item in payload["items"]},
|
|
{"203.0.113.9", "198.51.100.77"})
|
|
self.assertTrue(all(item["sources"] == "feature"
|
|
for item in payload["items"]))
|
|
|
|
with urllib.request.urlopen(base + "/api/stats", timeout=5) as response:
|
|
stats = json.loads(response.read().decode("utf-8"))
|
|
self.assertEqual(stats["core"], 2)
|
|
self.assertNotIn("review", stats)
|
|
|
|
with urllib.request.urlopen(base + "/api/ip/203.0.113.9", timeout=5) as response:
|
|
detail = json.loads(response.read().decode("utf-8"))
|
|
self.assertEqual(detail["item"]["tier"], watch.TIER_MALICIOUS)
|
|
self.assertFalse(detail["evidence"]["available"])
|
|
|
|
with urllib.request.urlopen(base + "/api/export.csv", timeout=5) as response:
|
|
body = response.read().decode("utf-8")
|
|
self.assertIn("text/csv", response.headers.get("Content-Type", ""))
|
|
self.assertIn("ip,tier,score", body)
|
|
self.assertIn("203.0.113.9", body)
|
|
self.assertIn("198.51.100.77", body)
|
|
|
|
with urllib.request.urlopen(base + "/api/export.txt", timeout=5) as response:
|
|
text = response.read().decode("utf-8")
|
|
self.assertEqual(set(text.split()),
|
|
{"203.0.113.9", "198.51.100.77"})
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
store.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|