200 lines
7.7 KiB
Python
200 lines
7.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Tests for the SQLite repositories + schema migrations (T001).
|
|
|
|
Run with:
|
|
python -m unittest discover -s tests -p 'test_*.py'
|
|
or (if pytest is installed):
|
|
python -m pytest tests/
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from app.db import connect
|
|
from app.db import repositories as repo
|
|
from app.db.schema import SCHEMA_VERSION, apply_migrations
|
|
|
|
|
|
def _new_conn() -> sqlite3.Connection:
|
|
fd, path = tempfile.mkstemp(suffix=".db")
|
|
os.close(fd)
|
|
os.unlink(path)
|
|
return connect(path)
|
|
|
|
|
|
class SchemaTest(unittest.TestCase):
|
|
def test_migrations_apply_and_are_idempotent(self):
|
|
conn = _new_conn()
|
|
self.assertEqual(conn.execute("PRAGMA user_version").fetchone()[0],
|
|
SCHEMA_VERSION)
|
|
# Re-applying should be a no-op (no exception, version unchanged).
|
|
apply_migrations(conn)
|
|
self.assertEqual(conn.execute("PRAGMA user_version").fetchone()[0],
|
|
SCHEMA_VERSION)
|
|
conn.close()
|
|
|
|
def test_core_tables_exist(self):
|
|
conn = _new_conn()
|
|
tables = {r["name"] for r in conn.execute(
|
|
"SELECT name FROM sqlite_master WHERE type='table'")}
|
|
for t in ("requests", "ip_stats", "scans", "scans_sources",
|
|
"blocklist_entries", "blocklist_events",
|
|
"blocklist_reasons", "whitelist_entries",
|
|
"whitelist_sources", "firewall_snapshots",
|
|
"firewall_address_list_entries", "firewall_rule_entries",
|
|
"whois_info", "log_sources"):
|
|
self.assertIn(t, tables, f"missing table {t}")
|
|
conn.close()
|
|
|
|
|
|
class ScanAndRequestTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.conn = _new_conn()
|
|
|
|
def tearDown(self):
|
|
self.conn.close()
|
|
|
|
def test_scan_lifecycle(self):
|
|
sid = repo.start_scan(self.conn, "ids", params={"full": True})
|
|
self.assertEqual(repo.get_scan(self.conn, sid)["status"], "running")
|
|
repo.finish_scan(self.conn, sid, status="done",
|
|
result={"ok": True}, request_count=3)
|
|
s = repo.get_scan(self.conn, sid)
|
|
self.assertEqual(s["status"], "done")
|
|
self.assertEqual(s["request_count"], 3)
|
|
self.assertEqual([x["id"] for x in repo.latest_scans(self.conn)],
|
|
[sid])
|
|
|
|
def test_insert_and_query_requests(self):
|
|
sid = repo.start_scan(self.conn, "ids")
|
|
rows = [
|
|
{"source": "npm", "scan_id": sid, "log_file": "a.log",
|
|
"ts": "2026-09-01 10:00:00", "ip": "1.2.3.4", "version": 4,
|
|
"method": "GET", "path": "/.env", "status": 404,
|
|
"host": "h", "user": None, "ua": "curl", "referer": None,
|
|
"denied": False, "banned": False, "is_error": True, "extra": None},
|
|
{"source": "copyparty", "scan_id": sid, "log_file": "f.txt",
|
|
"ts": "2026-09-01 10:00:01", "ip": "1.2.3.4", "version": 4,
|
|
"method": "GET", "path": "/admin", "status": 403,
|
|
"host": None, "user": None, "ua": None, "referer": None,
|
|
"denied": True, "banned": False, "is_error": True, "extra": None},
|
|
]
|
|
self.assertEqual(repo.insert_requests(self.conn, rows), 2)
|
|
self.assertEqual(repo.count_requests(self.conn), 2)
|
|
# cross-source
|
|
cs = repo.correlate_across_sources(self.conn)
|
|
self.assertEqual(cs[0]["ip"], "1.2.3.4")
|
|
self.assertEqual(cs[0]["sources"], 2)
|
|
# timeline
|
|
tl = repo.ip_timeline(self.conn, "1.2.3.4")
|
|
self.assertEqual(len(tl), 2)
|
|
# top ips
|
|
top = repo.top_ips(self.conn)
|
|
self.assertEqual(top[0]["total"], 2)
|
|
repo.finish_scan(self.conn, sid, status="done")
|
|
|
|
|
|
class BlocklistTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.conn = _new_conn()
|
|
|
|
def tearDown(self):
|
|
self.conn.close()
|
|
|
|
def test_grow_only_record(self):
|
|
r1 = repo.record_blocklist(self.conn, ["5.5.5.5/32", "6.6.6.6/32"],
|
|
added_by="scan:1")
|
|
self.assertEqual(r1["active"], 2)
|
|
# a later scan only sees 5.5.5.5; 6.6.6.6 stays (cumulative).
|
|
r2 = repo.record_blocklist(
|
|
self.conn, ["5.5.5.5/32"], added_by="scan:2")
|
|
active = [e["cidr"] for e in repo.active_blocklist(self.conn)]
|
|
self.assertIn("6.6.6.6/32", active)
|
|
self.assertEqual(r2["active"], 2)
|
|
|
|
def test_unblock(self):
|
|
repo.record_blocklist(self.conn, ["7.7.7.7/32"])
|
|
self.assertTrue(repo.unblock_entry(self.conn, "7.7.7.7/32"))
|
|
self.assertNotIn("7.7.7.7/32",
|
|
[e["cidr"] for e in repo.active_blocklist(self.conn)])
|
|
self.assertFalse(repo.unblock_entry(self.conn, "7.7.7.7/32"))
|
|
|
|
def test_prune_old_single_ip_blocks_keeps_ranges(self):
|
|
old = (datetime.now(timezone.utc) - timedelta(days=31)).isoformat(
|
|
timespec="seconds")
|
|
with repo.transaction(self.conn):
|
|
repo._block_entry_inner(self.conn, "1.2.3.4/32", last_seen=old)
|
|
repo._block_entry_inner(
|
|
self.conn, "2001:db8::1/128", last_seen=old)
|
|
repo._block_entry_inner(self.conn, "1.2.3.0/24", last_seen=old)
|
|
repo._block_entry_inner(self.conn, "2001:db8::/64", last_seen=old)
|
|
|
|
removed = repo.prune_old_single_ip_blocks(
|
|
self.conn, older_than_days=30)
|
|
self.assertCountEqual(removed, ["1.2.3.4/32", "2001:db8::1/128"])
|
|
|
|
active = {e["cidr"] for e in repo.active_blocklist(self.conn)}
|
|
self.assertIn("1.2.3.0/24", active)
|
|
self.assertIn("2001:db8::/64", active)
|
|
self.assertNotIn("1.2.3.4/32", active)
|
|
self.assertNotIn("2001:db8::1/128", active)
|
|
|
|
def test_prune_old_single_ip_blocks_respects_recency(self):
|
|
fresh = (datetime.now(timezone.utc) - timedelta(days=5)).isoformat(
|
|
timespec="seconds")
|
|
with repo.transaction(self.conn):
|
|
repo._block_entry_inner(self.conn, "9.9.9.9/32", last_seen=fresh)
|
|
|
|
removed = repo.prune_old_single_ip_blocks(
|
|
self.conn, older_than_days=30)
|
|
self.assertEqual(removed, [])
|
|
self.assertIn("9.9.9.9/32",
|
|
{e["cidr"] for e in repo.active_blocklist(self.conn)})
|
|
|
|
|
|
class WhitelistTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.conn = _new_conn()
|
|
|
|
def tearDown(self):
|
|
self.conn.close()
|
|
|
|
def test_whitelist_lookup(self):
|
|
repo.set_whitelist(self.conn, "Cloudflare", ["173.245.48.0/20"],
|
|
source="test")
|
|
self.assertTrue(repo.is_whitelisted_db(self.conn, "173.245.48.5"))
|
|
self.assertFalse(repo.is_whitelisted_db(self.conn, "8.8.8.8"))
|
|
# provider replace deactivates old ranges
|
|
repo.set_whitelist(self.conn, "Cloudflare", ["104.16.0.0/13"])
|
|
self.assertFalse(repo.is_whitelisted_db(self.conn, "173.245.48.5"))
|
|
|
|
|
|
class FirewallTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.conn = _new_conn()
|
|
|
|
def tearDown(self):
|
|
self.conn.close()
|
|
|
|
def test_firewall_snapshot(self):
|
|
sid = repo.save_firewall_snapshot(
|
|
self.conn,
|
|
address_lists_v4=[{"list": "Blocked", "address": "5.5.5.5/32"}],
|
|
address_lists_v6=[{"list": "Init7", "address": "2001:1620::/30"}],
|
|
rules_v4=[{"chain": "input", "action": "drop", "comment": "x"}],
|
|
rules_v6=[])
|
|
self.assertEqual(sid, 1)
|
|
self.assertEqual(self.conn.execute(
|
|
"SELECT COUNT(*) FROM firewall_address_list_entries").fetchone()[0], 2)
|
|
self.assertEqual(self.conn.execute(
|
|
"SELECT COUNT(*) FROM firewall_rule_entries").fetchone()[0], 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|