fix: enhance logging and improve database insert handling 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:26:37 +01:00
parent 18a9b59d96
commit c762942f70

View File

@@ -1,31 +1,65 @@
import asyncio import asyncio
from scapy.all import sniff, Ether, IP from scapy.all import sniff, Ether, IP
import asyncpg import asyncpg
from typing import List
import threading import threading
import logging
from typing import List, Dict
# ------------------------------------------------------------------------------
# LOGGING
# ------------------------------------------------------------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("sniffer")
# ------------------------------------------------------------------------------
# CONFIG
# ------------------------------------------------------------------------------
DB_DSN = "postgresql://mitm_user:mitm_password@localhost:5432/mitm_db" DB_DSN = "postgresql://mitm_user:mitm_password@localhost:5432/mitm_db"
# interface → direction map
# Example: { "eth0": "ingress", "eth1": "egress" }
iface_directions: Dict[str, str] = {}
sniffer_threads = {} sniffer_threads = {}
event_loop = None # <-- we store FastAPI's main loop here
# ------------------------------------------------------------------------------
# ASYNCIO LOOP (for DB INSERTS)
# ------------------------------------------------------------------------------
async_loop = asyncio.new_event_loop()
# ----------------------------- def start_async_loop(loop):
asyncio.set_event_loop(loop)
loop.run_forever()
threading.Thread(
target=start_async_loop,
args=(async_loop,),
daemon=True
).start()
logger.info("Asyncio DB event loop started")
# ------------------------------------------------------------------------------
# DB INSERT # DB INSERT
# ----------------------------- # ------------------------------------------------------------------------------
async def db_insert_packet(pkt_info: dict): async def db_insert_packet(pkt_info: dict):
conn = await asyncpg.connect(DB_DSN) """
Insert a packet asynchronously into PostgreSQL.
"""
try: try:
await conn.execute(""" conn = await asyncpg.connect(DB_DSN)
INSERT INTO packets( await conn.execute(
iface, direction, """
INSERT INTO packet_log(
interface, direction,
src_mac, dst_mac, eth_type, src_mac, dst_mac, eth_type,
src_ip, dst_ip, protocol, src_ip, dst_ip, ip_protocol,
length, ebpf_verdict, ebpf_chain, raw packet_len, nft_verdict, nft_chain, raw_packet
) )
VALUES( VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)
$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12
);
""", """,
pkt_info["iface"], pkt_info["iface"],
pkt_info["direction"], pkt_info["direction"],
@@ -36,23 +70,45 @@ async def db_insert_packet(pkt_info: dict):
pkt_info["dst_ip"], pkt_info["dst_ip"],
pkt_info["protocol"], pkt_info["protocol"],
pkt_info["length"], pkt_info["length"],
pkt_info["ebpf_verdict"], pkt_info.get("ebpf_verdict"),
pkt_info["ebpf_chain"], pkt_info.get("ebpf_chain"),
pkt_info["raw"] pkt_info["raw"]
) )
logger.info(
f"DB Insert OK [{pkt_info['iface']}/{pkt_info['direction']}] "
f"{pkt_info['src_ip']} → {pkt_info['dst_ip']} len={pkt_info['length']}"
)
except Exception as e:
logger.exception(f"DB insert failed: {e}")
finally: finally:
try:
await conn.close() await conn.close()
except:
pass
# ------------------------------------------------------------------------------
# PACKET HANDLER
# ------------------------------------------------------------------------------
def detect_direction(iface: str) -> str:
"""
Return ingress/egress direction for this interface.
"""
return iface_directions.get(iface, "unknown")
# -----------------------------
# PACKET HANDLER (Thread)
# -----------------------------
def handle_packet(pkt, iface): def handle_packet(pkt, iface):
global event_loop """
Called inside Scapy sniff thread.
"""
direction = detect_direction(iface)
pkt_info = { pkt_info = {
"iface": iface, "iface": iface,
"direction": "unknown", "direction": direction,
"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,
@@ -62,44 +118,63 @@ def handle_packet(pkt, iface):
"length": len(pkt), "length": len(pkt),
"ebpf_verdict": None, "ebpf_verdict": None,
"ebpf_chain": None, "ebpf_chain": None,
"raw": bytes(pkt), "raw": bytes(pkt)
} }
# Schedule coroutine in main event loop from thread # Debug per-packet
asyncio.run_coroutine_threadsafe( logger.debug(
db_insert_packet(pkt_info), f"Packet captured on {iface} ({direction}): "
event_loop f"{pkt_info['src_ip']} -> {pkt_info['dst_ip']} ({pkt_info['length']} bytes)"
) )
# Submit to asyncio loop
asyncio.run_coroutine_threadsafe(db_insert_packet(pkt_info), async_loop)
# ----------------------------- # ------------------------------------------------------------------------------
# SNIFFER THREAD # 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), logger.info(f"Starting sniffer on {iface}")
iface=iface, sniff(prn=lambda x: handle_packet(x, iface), iface=iface, store=False)
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 # PUBLIC API
# ----------------------------- # ------------------------------------------------------------------------------
async def start_sniffing(interfaces: List[str]):
global event_loop async def start_sniffing(interfaces: List[str], ingress_iface: str, egress_iface: str):
event_loop = asyncio.get_running_loop() # <-- IMPORTANT FIX """
Start sniffing on a list of interfaces.
Also store ingress/egress direction info.
"""
# Save direction mapping
iface_directions.clear()
iface_directions[ingress_iface] = "ingress"
iface_directions[egress_iface] = "egress"
logger.info(f"Direction mapping: {iface_directions}")
for iface in interfaces: for iface in interfaces:
if iface in sniffer_threads: if iface in sniffer_threads:
logger.info(f"Sniffer already active on {iface}")
continue continue
sniffer_threads[iface] = start_sniffer_thread(iface) sniffer_threads[iface] = start_sniffer_thread(iface)
return {"status": "sniffers started", "interfaces": interfaces}
async def stop_sniffing(): async def stop_sniffing():
# Scapy cannot stop sniff(), but threads die on shutdown """
Scapy sniff cannot stop easily.
We simply forget threads (they are daemon threads).
"""
logger.warning("Stopping all sniffers (threads will exit on process stop)")
sniffer_threads.clear() sniffer_threads.clear()
return True return True