131 lines
6.0 KiB
Python
131 lines
6.0 KiB
Python
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, os.path.dirname(HERE))
|
|
|
|
from graylog_ops import watch # noqa: E402
|
|
|
|
NOW = "2026-10-17T16:00:00Z"
|
|
|
|
|
|
def _behavior(ip, verdict="建议封禁", score=75):
|
|
return {
|
|
"ip": ip, "score": score, "verdict": verdict,
|
|
"signals": ["高失败率"], "reason_text": "24 小时内扫描了大量路径。",
|
|
"share_note": "", "behavior_type": "目录扫描",
|
|
"asn": "4538", "asn_org": "教育网", "shared_type": "education",
|
|
"total": 50, "pct4": 60.0, "uri_card": 20, "uri_templates": 15,
|
|
"ua_card": 1, "host_card": 2, "coverage": 100.0, "peak_share": 50.0,
|
|
"top_uri": "/.env", "top_ua": "curl/8.0",
|
|
"country_code": "CN", "country_name": "中国", "region": "北京",
|
|
"city": "Beijing",
|
|
}
|
|
|
|
|
|
class FusedTests(unittest.TestCase):
|
|
def test_feature_and_behavior_merge(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
try:
|
|
store.upsert({"ip": "203.0.113.9", "tier": watch.TIER_MALICIOUS,
|
|
"score": 91.0, "total": 120, "window": 3600,
|
|
"rules": {"probe_java": 10},
|
|
"reasons": ["探测 10 次"]}, NOW)
|
|
store.sync_fused(NOW)
|
|
fused = store.fused_items()
|
|
self.assertEqual(len(fused), 1)
|
|
self.assertEqual(fused[0]["sources"], "feature")
|
|
self.assertEqual(fused[0]["first_banned_at"], NOW)
|
|
self.assertIn("特征命中", fused[0]["fused_reason"])
|
|
|
|
# behaviour-only ban is merged
|
|
store.replace_behavior_profiles([_behavior("198.51.100.9")])
|
|
store.sync_fused("2026-10-17T16:10:00Z")
|
|
merged = {item["ip"]: item for item in store.fused_items()}
|
|
self.assertIn("198.51.100.9", merged)
|
|
self.assertEqual(merged["198.51.100.9"]["sources"], "behavior")
|
|
self.assertEqual(merged["203.0.113.9"]["sources"], "feature")
|
|
|
|
# same IP on both sides becomes "both"
|
|
store.upsert({"ip": "198.51.100.9", "tier": watch.TIER_MALICIOUS,
|
|
"score": 80.0, "total": 50, "window": 3600,
|
|
"rules": {"probe_java": 6},
|
|
"reasons": ["探测 6 次"]}, "2026-10-17T16:11:00Z")
|
|
store.sync_fused("2026-10-17T16:12:00Z")
|
|
merged = {item["ip"]: item for item in store.fused_items()}
|
|
self.assertEqual(merged["198.51.100.9"]["sources"], "both")
|
|
|
|
# watch-only behaviour must not enter the block list
|
|
store.replace_behavior_profiles([
|
|
_behavior("192.0.2.44", verdict="仅告警（共享出口）")])
|
|
store.sync_fused("2026-10-17T16:20:00Z")
|
|
ips = {item["ip"] for item in store.fused_items()}
|
|
self.assertNotIn("192.0.2.44", ips)
|
|
self.assertIn("203.0.113.9", ips) # feature rows stay forever
|
|
finally:
|
|
store.close()
|
|
|
|
def test_push_throttle_24h(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
try:
|
|
now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
|
store.replace_behavior_profiles([_behavior("198.51.100.7")])
|
|
store.sync_fused(now)
|
|
# first sync suppresses historical rows; a fresh detection is due
|
|
with store.lock, store.conn:
|
|
store.conn.execute(
|
|
"UPDATE ip_fused SET last_pushed_at=NULL WHERE ip=?",
|
|
("198.51.100.7",))
|
|
due = store.due_pushes(active_hours=24, cooldown_hours=24)
|
|
self.assertEqual([row["ip"] for row in due], ["198.51.100.7"])
|
|
|
|
store.mark_pushed(["198.51.100.7"], now)
|
|
self.assertEqual(store.due_pushes(24, 24), []) # throttled
|
|
|
|
old = time.strftime("%Y-%m-%dT%H:%M:%SZ",
|
|
time.gmtime(time.time() - 25 * 3600))
|
|
with store.lock, store.conn:
|
|
store.conn.execute(
|
|
"UPDATE ip_fused SET last_pushed_at=? WHERE ip=?",
|
|
(old, "198.51.100.7"))
|
|
due = store.due_pushes(active_hours=24, cooldown_hours=24)
|
|
self.assertEqual([row["ip"] for row in due], ["198.51.100.7"])
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
class HistoricalBehaviourTests(unittest.TestCase):
|
|
def test_historical_profile_survives_live_refresh(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store = watch.Store(os.path.join(directory, "ips.db"))
|
|
try:
|
|
store.upsert_historical_profiles([{
|
|
"ip": "198.51.100.9", "score": 80,
|
|
"verdict": "建议封禁", "signals": ["历史补测"],
|
|
"last_ts": "2026-10-18T10:00:00Z", "total": 1000,
|
|
}])
|
|
store.replace_behavior_profiles([{
|
|
"ip": "203.0.113.9", "score": 70,
|
|
"verdict": "仅告警（共享出口）"}])
|
|
rows = {r["ip"]: dict(r) for r in store.conn.execute(
|
|
"SELECT * FROM behavior_profiles")}
|
|
self.assertEqual(rows["198.51.100.9"]["historical"], 1)
|
|
self.assertEqual(rows["203.0.113.9"]["historical"], 0)
|
|
|
|
store.sync_fused()
|
|
fused = {r["ip"]: dict(r) for r in store.conn.execute(
|
|
"SELECT * FROM ip_fused")}
|
|
self.assertEqual(fused["198.51.100.9"]["first_banned_at"],
|
|
"2026-10-18T10:00:00Z")
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|