fix: improve event loop handling and code organization in network sniffer
All checks were successful
Build and Deploy MITM Webserver / build (push) Successful in 7s

This commit is contained in:
2025-11-25 16:22:03 +01:00
parent bbdef77e0e
commit 18a9b59d96

View File

@@ -6,9 +6,13 @@ import threading
DB_DSN = "postgresql://mitm_user:mitm_password@localhost:5432/mitm_db" DB_DSN = "postgresql://mitm_user:mitm_password@localhost:5432/mitm_db"
sniffer_tasks = {}
sniffer_threads = {} sniffer_threads = {}
event_loop = None # <-- we store FastAPI's main loop here
# -----------------------------
# DB INSERT
# -----------------------------
async def db_insert_packet(pkt_info: dict): async def db_insert_packet(pkt_info: dict):
conn = await asyncpg.connect(DB_DSN) conn = await asyncpg.connect(DB_DSN)
try: try:
@@ -40,13 +44,15 @@ async def db_insert_packet(pkt_info: dict):
await conn.close() await conn.close()
# -----------------------------
# PACKET HANDLER (Thread)
# -----------------------------
def handle_packet(pkt, iface): def handle_packet(pkt, iface):
""" global event_loop
Scapy callback → run in thread.
"""
pkt_info = { pkt_info = {
"iface": iface, "iface": iface,
"direction": "unknown", # You can set eth0=ingress, eth1=egress "direction": "unknown",
"src_mac": pkt[Ether].src if Ether in pkt else None, "src_mac": pkt[Ether].src if Ether in pkt else None,
"dst_mac": pkt[Ether].dst if Ether in pkt else None, "dst_mac": pkt[Ether].dst if Ether in pkt else None,
"eth_type": pkt[Ether].type if Ether in pkt else None, "eth_type": pkt[Ether].type if Ether in pkt else None,
@@ -54,26 +60,39 @@ def handle_packet(pkt, iface):
"dst_ip": pkt[IP].dst if IP in pkt else None, "dst_ip": pkt[IP].dst if IP in pkt else None,
"protocol": pkt[IP].proto if IP in pkt else None, "protocol": pkt[IP].proto if IP in pkt else None,
"length": len(pkt), "length": len(pkt),
"ebpf_verdict": None, # will integrate later "ebpf_verdict": None,
"ebpf_chain": None, "ebpf_chain": None,
"raw": bytes(pkt) "raw": bytes(pkt),
} }
# push to asyncio loop # Schedule coroutine in main event loop from thread
loop = asyncio.get_event_loop() asyncio.run_coroutine_threadsafe(
loop.create_task(db_insert_packet(pkt_info)) db_insert_packet(pkt_info),
event_loop
)
# -----------------------------
# SNIFFER THREAD
# -----------------------------
def start_sniffer_thread(iface: str): def start_sniffer_thread(iface: str):
def sniff_blocking(): def sniff_blocking():
sniff(prn=lambda x: handle_packet(x, iface), iface=iface, store=False) sniff(prn=lambda x: handle_packet(x, iface),
iface=iface,
store=False)
thread = threading.Thread(target=sniff_blocking, daemon=True) thread = threading.Thread(target=sniff_blocking, daemon=True)
thread.start() thread.start()
return thread return thread
# -----------------------------
# PUBLIC API
# -----------------------------
async def start_sniffing(interfaces: List[str]): async def start_sniffing(interfaces: List[str]):
global event_loop
event_loop = asyncio.get_running_loop() # <-- IMPORTANT FIX
for iface in interfaces: for iface in interfaces:
if iface in sniffer_threads: if iface in sniffer_threads:
continue continue
@@ -81,6 +100,6 @@ async def start_sniffing(interfaces: List[str]):
async def stop_sniffing(): async def stop_sniffing():
# Scapy cannot easily stop sniff(), so we simply kill threads # Scapy cannot stop sniff(), but threads die on shutdown
sniffer_threads.clear() sniffer_threads.clear()
return True return True