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
All checks were successful
Build and Deploy MITM Webserver / build (push) Successful in 7s
This commit is contained in:
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user