#!/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()