170 lines
7.6 KiB
Python
170 lines
7.6 KiB
Python
import os
|
|
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
|
|
from app.store import AlertStore
|
|
|
|
|
|
class StoreTests(unittest.TestCase):
|
|
def test_insert_and_read(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
path = os.path.join(td, "alerts.db")
|
|
store = AlertStore(path)
|
|
event = {
|
|
"timestamp": "2026-08-13T10:00:00+00:00",
|
|
"event_type": "alert",
|
|
"src_ip": "192.168.100.10",
|
|
"src_port": 12345,
|
|
"dest_ip": "9.9.9.9",
|
|
"dest_port": 443,
|
|
"proto": "TCP",
|
|
"alert": {
|
|
"signature_id": 1000001,
|
|
"signature": "LOCAL TZSP PIPELINE TEST",
|
|
"category": "Test",
|
|
"severity": 1,
|
|
"action": "allowed",
|
|
},
|
|
}
|
|
store.insert_alert(event, False, "9.9.9.9", "observation")
|
|
rows = store.recent(10)
|
|
self.assertEqual(len(rows), 1)
|
|
self.assertEqual(rows[0]["signature_id"], 1000001)
|
|
self.assertEqual(store.summary()["total_alerts"], 1)
|
|
store.close()
|
|
|
|
def test_deduplicates_and_aggregates_hits(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
path = os.path.join(td, "alerts.db")
|
|
store = AlertStore(path)
|
|
event = {
|
|
"timestamp": "2099-08-13T10:00:00+00:00",
|
|
"event_type": "alert",
|
|
"src_ip": "192.168.100.10",
|
|
"src_port": 12345,
|
|
"dest_ip": "9.9.9.9",
|
|
"dest_port": 443,
|
|
"proto": "TCP",
|
|
"alert": {"signature_id": 42, "signature": "duplicate", "severity": 1},
|
|
}
|
|
alert_id = store.insert_alert(event, False, "9.9.9.9", "observation")
|
|
self.assertEqual(store.find_recent_duplicate(event, 300), alert_id)
|
|
store.bump_duplicate(alert_id, event)
|
|
row = store.recent(1)[0]
|
|
self.assertEqual(row["hit_count"], 2)
|
|
self.assertEqual(store.summary()["total_alerts"], 2)
|
|
self.assertEqual(store.summary()["incidents"], 1)
|
|
store.close()
|
|
|
|
def test_migrates_pre_040_database(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
path = os.path.join(td, "alerts.db")
|
|
conn = sqlite3.connect(path)
|
|
conn.executescript(
|
|
"""
|
|
CREATE TABLE alerts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp TEXT NOT NULL,
|
|
flow_id TEXT, src_ip TEXT, src_port INTEGER,
|
|
dest_ip TEXT, dest_port INTEGER, proto TEXT,
|
|
signature_id INTEGER, signature TEXT, category TEXT,
|
|
severity INTEGER, action TEXT, blocked INTEGER NOT NULL DEFAULT 0,
|
|
block_target TEXT, block_reason TEXT, raw_json TEXT NOT NULL
|
|
);
|
|
INSERT INTO alerts (timestamp, signature_id, signature, blocked, raw_json)
|
|
VALUES ('2026-08-13T10:00:00+00:00', 7, 'legacy', 0, '{}');
|
|
"""
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
store = AlertStore(path)
|
|
row = store.recent(1)[0]
|
|
self.assertEqual(row["hit_count"], 1)
|
|
self.assertEqual(row["first_seen"], "2026-08-13T10:00:00+00:00")
|
|
self.assertEqual(store.database_info()["schema_version"], 4)
|
|
store.close()
|
|
|
|
def test_normalizes_timezone_to_utc(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
path = os.path.join(td, "alerts.db")
|
|
store = AlertStore(path)
|
|
event = {
|
|
"timestamp": "2026-08-14T10:14:21+02:00",
|
|
"src_ip": "192.0.2.10",
|
|
"dest_ip": "198.51.100.20",
|
|
"dest_port": 443,
|
|
"proto": "TCP",
|
|
"alert": {"signature_id": 77, "signature": "timezone", "severity": 2},
|
|
}
|
|
store.insert_alert(event, False, None, "observation")
|
|
self.assertEqual(store.recent(1)[0]["last_seen"], "2026-08-14T08:14:21+00:00")
|
|
store.close()
|
|
|
|
def test_upgrade_compacts_legacy_duplicate_incidents(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
path = os.path.join(td, "alerts.db")
|
|
conn = sqlite3.connect(path)
|
|
conn.executescript(
|
|
"""
|
|
CREATE TABLE alerts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp TEXT NOT NULL, first_seen TEXT, last_seen TEXT,
|
|
hit_count INTEGER NOT NULL DEFAULT 1, flow_id TEXT,
|
|
src_ip TEXT, src_port INTEGER, dest_ip TEXT, dest_port INTEGER, proto TEXT,
|
|
signature_id INTEGER, signature TEXT, category TEXT, severity INTEGER, action TEXT,
|
|
blocked INTEGER NOT NULL DEFAULT 0, block_target TEXT, block_reason TEXT, raw_json TEXT NOT NULL
|
|
);
|
|
PRAGMA user_version=3;
|
|
INSERT INTO alerts (timestamp,first_seen,last_seen,src_ip,dest_ip,dest_port,proto,signature_id,signature,severity,raw_json)
|
|
VALUES ('2026-08-14T10:14:20+02:00','2026-08-14T10:14:20+02:00','2026-08-14T10:14:20+02:00','109.173.161.12','1.1.1.1',NULL,'ICMP',42,'legacy duplicate',2,'{}');
|
|
INSERT INTO alerts (timestamp,first_seen,last_seen,src_ip,dest_ip,dest_port,proto,signature_id,signature,severity,raw_json)
|
|
VALUES ('2026-08-14T08:16:12+00:00','2026-08-14T08:16:12+00:00','2026-08-14T08:16:12+00:00','109.173.161.12','1.1.1.1',NULL,'ICMP',42,'legacy duplicate',2,'{}');
|
|
"""
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
store = AlertStore(path)
|
|
rows = store.recent(10)
|
|
self.assertEqual(len(rows), 1)
|
|
self.assertEqual(rows[0]["hit_count"], 2)
|
|
self.assertEqual(rows[0]["first_seen"], "2026-08-14T08:14:20+00:00")
|
|
self.assertEqual(rows[0]["last_seen"], "2026-08-14T08:16:12+00:00")
|
|
store.close()
|
|
|
|
def test_purges_reserved_selftest_sid(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
path = os.path.join(td, "alerts.db")
|
|
store = AlertStore(path)
|
|
event = {
|
|
"timestamp": "2026-08-14T08:00:00+00:00",
|
|
"alert": {"signature_id": 1000001, "signature": "old test name", "severity": 3},
|
|
}
|
|
store.insert_alert(event, False, None, "selftest")
|
|
self.assertEqual(store.purge_builtin_test_incidents(), 1)
|
|
self.assertEqual(store.summary()["incidents"], 0)
|
|
store.close()
|
|
|
|
def test_incident_window_is_bounded_from_first_seen(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
path = os.path.join(td, "alerts.db")
|
|
store = AlertStore(path)
|
|
base = {
|
|
"src_ip": "192.0.2.1", "dest_ip": "198.51.100.1",
|
|
"dest_port": 22, "proto": "TCP",
|
|
"alert": {"signature_id": 88, "signature": "window", "severity": 1},
|
|
}
|
|
first = dict(base, timestamp="2026-08-14T08:00:00+00:00")
|
|
middle = dict(base, timestamp="2026-08-14T08:04:59+00:00")
|
|
later = dict(base, timestamp="2026-08-14T08:05:01+00:00")
|
|
alert_id = store.insert_alert(first, False, None, "observation")
|
|
self.assertEqual(store.find_recent_duplicate(middle, 300), alert_id)
|
|
store.bump_duplicate(alert_id, middle)
|
|
self.assertIsNone(store.find_recent_duplicate(later, 300))
|
|
store.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|