poc2_worked

tests/test_flow_tracker.py

@@ -0,0 +1,71 @@
+import socket
+import struct
+import time
+import unittest
+
+from app.flow_tracker import FlowTracker, _parse_frame
+from app.live import TrafficNormalizer
+
+
+class _Pipeline:
+    def __init__(self):
+        self.rows = []
+
+    def publish(self, event, persist=True):
+        self.rows.append((event, persist))
+
+
+def _ipv4_tcp_frame(src: str, sport: int, dst: str, dport: int, payload: bytes = b"") -> bytes:
+    eth = b"\x00" * 12 + struct.pack("!H", 0x0800)
+    total_len = 20 + 20 + len(payload)
+    ip = struct.pack(
+        "!BBHHHBBH4s4s",
+        0x45,
+        0,
+        total_len,
+        1,
+        0,
+        64,
+        6,
+        0,
+        socket.inet_aton(src),
+        socket.inet_aton(dst),
+    )
+    tcp = struct.pack("!HHLLBBHHH", sport, dport, 0, 0, 5 << 4, 0x10, 65535, 0, 0)
+    return eth + ip + tcp + payload
+
+
+class FlowTrackerTests(unittest.TestCase):
+    def test_parses_ipv4_tcp_tuple(self):
+        frame = _ipv4_tcp_frame("192.168.100.10", 51000, "1.1.1.1", 443)
+        self.assertEqual(_parse_frame(frame), ("192.168.100.10", 51000, "1.1.1.1", 443, "TCP"))
+
+    def test_reverse_packets_update_one_live_session_without_persistence(self):
+        pipeline = _Pipeline()
+        tracker = FlowTracker(
+            TrafficNormalizer("192.168.100.0/24"),
+            pipeline,  # type: ignore[arg-type]
+            update_interval_seconds=0.25,
+            max_flows=1000,
+        )
+        outbound = _ipv4_tcp_frame("192.168.100.10", 51000, "1.1.1.1", 443, b"hello")
+        inbound = _ipv4_tcp_frame("1.1.1.1", 443, "192.168.100.10", 51000, b"world")
+
+        tracker.observe(outbound)
+        time.sleep(0.26)
+        tracker.observe(inbound)
+
+        self.assertEqual(tracker.status()["active_flows"], 1)
+        self.assertEqual(len(pipeline.rows), 2)
+        first, first_persist = pipeline.rows[0]
+        second, second_persist = pipeline.rows[1]
+        self.assertEqual(first["id"], second["id"])
+        self.assertEqual(second["direction"], "outbound")
+        self.assertEqual(second["app_proto"], "tls")
+        self.assertGreater(second["bytes"], first["bytes"])
+        self.assertFalse(first_persist)
+        self.assertFalse(second_persist)
+
+
+if __name__ == "__main__":
+    unittest.main()