98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
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)
|
|
|
|
def test_status_exposes_monotonic_directional_traffic_counters(self):
|
|
pipeline = _Pipeline()
|
|
tracker = FlowTracker(
|
|
TrafficNormalizer("192.168.100.0/24"),
|
|
pipeline, # type: ignore[arg-type]
|
|
max_flows=1000,
|
|
)
|
|
outbound = _ipv4_tcp_frame("192.168.100.10", 51000, "1.1.1.1", 443, b"out")
|
|
inbound = _ipv4_tcp_frame("1.1.1.1", 443, "192.168.100.10", 51000, b"in")
|
|
internal = _ipv4_tcp_frame("192.168.100.10", 51000, "192.168.100.20", 443, b"lan")
|
|
external = _ipv4_tcp_frame("1.1.1.1", 51000, "8.8.8.8", 443, b"wan")
|
|
|
|
for frame in (outbound, inbound, internal, external):
|
|
tracker.observe(frame)
|
|
|
|
traffic = tracker.status()["traffic_counters"]
|
|
self.assertEqual(traffic["packets_total"], 4)
|
|
self.assertEqual(traffic["packets_out"], 1)
|
|
self.assertEqual(traffic["packets_in"], 1)
|
|
self.assertEqual(traffic["packets_internal"], 1)
|
|
self.assertEqual(traffic["packets_external"], 1)
|
|
self.assertEqual(
|
|
traffic["bytes_total"],
|
|
traffic["bytes_out"] + traffic["bytes_in"] + traffic["bytes_internal"] + traffic["bytes_external"],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|