fix: enhance logging and improve database insert handling in network sniffer

backend/src/network_sniffer.py

@@ -1,31 +1,65 @@
 import asyncio
 from scapy.all import sniff, Ether, IP
 import asyncpg
-from typing import List
 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"
 
+# interface → direction map
+# Example: { "eth0": "ingress", "eth1": "egress" }
+iface_directions: Dict[str, str] = {}
+
 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
-# -----------------------------
+# ------------------------------------------------------------------------------
+
 async def db_insert_packet(pkt_info: dict):
-    conn = await asyncpg.connect(DB_DSN)
+    """
+    Insert a packet asynchronously into PostgreSQL.
+    """
     try:
-        await conn.execute("""
-            INSERT INTO packets(
-                iface, direction,
+        conn = await asyncpg.connect(DB_DSN)
+        await conn.execute(
+            """
+            INSERT INTO packet_log(
+                interface, direction,
                 src_mac, dst_mac, eth_type,
-                src_ip, dst_ip, protocol,
-                length, ebpf_verdict, ebpf_chain, raw
+                src_ip, dst_ip, ip_protocol,
+                packet_len, nft_verdict, nft_chain, raw_packet
             )
-            VALUES(
-                $1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12
-            );
+            VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)
             """,
             pkt_info["iface"],
             pkt_info["direction"],
@@ -36,23 +70,45 @@ async def db_insert_packet(pkt_info: dict):
             pkt_info["dst_ip"],
             pkt_info["protocol"],
             pkt_info["length"],
-        pkt_info["ebpf_verdict"],
-        pkt_info["ebpf_chain"],
+            pkt_info.get("ebpf_verdict"),
+            pkt_info.get("ebpf_chain"),
             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:
+        try:
             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):
-    global event_loop
+    """
+    Called inside Scapy sniff thread.
+    """
+    direction = detect_direction(iface)
 
     pkt_info = {
         "iface": iface,
-        "direction": "unknown",
+        "direction": direction,
         "src_mac": pkt[Ether].src 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,
@@ -62,44 +118,63 @@ def handle_packet(pkt, iface):
         "length": len(pkt),
         "ebpf_verdict": None,
         "ebpf_chain": None,
-        "raw": bytes(pkt),
+        "raw": bytes(pkt)
     }
 
-    # Schedule coroutine in main event loop from thread
-    asyncio.run_coroutine_threadsafe(
-        db_insert_packet(pkt_info),
-        event_loop
+    # Debug per-packet
+    logger.debug(
+        f"Packet captured on {iface} ({direction}): "
+        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
-# -----------------------------
+# ------------------------------------------------------------------------------
+
 def start_sniffer_thread(iface: str):
     def sniff_blocking():
-        sniff(prn=lambda x: handle_packet(x, iface),
-              iface=iface,
-              store=False)
+        logger.info(f"Starting sniffer on {iface}")
+        sniff(prn=lambda x: handle_packet(x, iface), iface=iface, store=False)
 
     thread = threading.Thread(target=sniff_blocking, daemon=True)
     thread.start()
     return thread
 
-
-# -----------------------------
+# ------------------------------------------------------------------------------
 # PUBLIC API
-# -----------------------------
-async def start_sniffing(interfaces: List[str]):
-    global event_loop
-    event_loop = asyncio.get_running_loop()   # <-- IMPORTANT FIX
+# ------------------------------------------------------------------------------
+
+async def start_sniffing(interfaces: List[str], ingress_iface: str, egress_iface: str):
+    """
+    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:
         if iface in sniffer_threads:
+            logger.info(f"Sniffer already active on {iface}")
             continue
+
         sniffer_threads[iface] = start_sniffer_thread(iface)
 
+    return {"status": "sniffers started", "interfaces": interfaces}
+
 
 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()
     return True