70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
|
import sqlite3
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
from netscaler_triage.correlation import correlate_database
|
||
|
|
|
||
|
|
|
||
|
|
class CorrelationTests(unittest.TestCase):
|
||
|
|
def setUp(self) -> None:
|
||
|
|
self.db = sqlite3.connect(":memory:")
|
||
|
|
self.db.executescript("""
|
||
|
|
CREATE TABLE events(id INTEGER PRIMARY KEY,appliance TEXT,source TEXT,
|
||
|
|
line INTEGER,timestamp TEXT,kind TEXT,data TEXT);
|
||
|
|
CREATE TABLE event_keys(event_id INTEGER,appliance TEXT,key_type TEXT,
|
||
|
|
key_value TEXT,timestamp TEXT);
|
||
|
|
""")
|
||
|
|
|
||
|
|
def tearDown(self) -> None:
|
||
|
|
self.db.close()
|
||
|
|
|
||
|
|
def add(
|
||
|
|
self,
|
||
|
|
number: int,
|
||
|
|
stamp: str | None,
|
||
|
|
source: str = "access",
|
||
|
|
device: str = "one",
|
||
|
|
key: str = "source_ip",
|
||
|
|
value: str = "192.0.2.1",
|
||
|
|
) -> None:
|
||
|
|
self.db.execute(
|
||
|
|
"INSERT INTO events VALUES(?,?,?,?,?,?,?)",
|
||
|
|
(number, device, source, number, stamp, "http", "{}"),
|
||
|
|
)
|
||
|
|
self.db.execute(
|
||
|
|
"INSERT INTO event_keys VALUES(?,?,?,?,?)", (number, device, key, value, stamp)
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_nearest_other_log_and_device_isolation(self) -> None:
|
||
|
|
self.add(1, "2026-01-01T00:00:00Z", "auth")
|
||
|
|
self.add(2, "2026-01-01T00:00:01Z")
|
||
|
|
self.add(3, "2026-01-01T00:00:02Z")
|
||
|
|
self.add(4, "2026-01-01T00:00:03Z", device="two")
|
||
|
|
result = correlate_database(self.db)
|
||
|
|
self.assertEqual(result["links"], 2)
|
||
|
|
self.assertEqual(
|
||
|
|
self.db.execute("SELECT left_id,right_id FROM correlations").fetchall(),
|
||
|
|
[(1, 2), (1, 3)],
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_times_limits_and_rerun(self) -> None:
|
||
|
|
self.add(1, "2026-01-01T01:00:00+01:00", "auth")
|
||
|
|
self.add(2, "2026-01-01T00:00:01Z")
|
||
|
|
self.add(3, "2026-01-01T00:00:02Z")
|
||
|
|
self.add(4, None)
|
||
|
|
self.add(5, "2026-01-01T00:00:03")
|
||
|
|
result = correlate_database(self.db, max_links=1)
|
||
|
|
self.assertTrue(result["truncated"])
|
||
|
|
result = correlate_database(self.db)
|
||
|
|
self.assertEqual(result["links"], 2)
|
||
|
|
self.assertEqual(result["invalid_time_keys"], 2)
|
||
|
|
self.assertEqual(correlate_database(self.db, window_seconds=0)["links"], 0)
|
||
|
|
|
||
|
|
def test_exact_keys_and_no_same_source_links(self) -> None:
|
||
|
|
self.add(1, "2026-01-01T00:00:00Z")
|
||
|
|
self.add(2, "2026-01-01T00:00:01Z")
|
||
|
|
self.add(3, "2026-01-01T00:00:02Z", "auth", value="192.0.2.10")
|
||
|
|
self.assertEqual(correlate_database(self.db)["links"], 0)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|