use shared objects like loops db access etc

backend/src/network_sniffer.py

@@ -1,3 +1,4 @@
+# src/sniffer.py
 import asyncio
 import logging
 import threading
@@ -8,8 +9,9 @@ import errno
 import struct
 from typing import Dict, List, Optional, Any
 
-import asyncpg
-from src.utilities.database import DB
+# NOTE: ensure this path points to your shared runtime module
+from backend.src import shared_objects
+
 from scapy.all import (
     Ether,
     ARP,
@@ -35,9 +37,6 @@ from src.Models.ip_protocol import IPProtocolEnum, protocol_from_number
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("af_packet_sniffer")
 
-# ---- Config -----------------------------------------------------------
-DB_DSN = "postgresql://mitm_user:mitm_password@localhost:5432/mitm_db"
-
 # ---- Globals (kept minimal) -------------------------------------------
 af_sockets: Dict[str, socket.socket] = {}
 af_thread: Optional[threading.Thread] = None
@@ -46,7 +45,11 @@ af_stop_event: Optional[threading.Event] = None
 fixed_bridge_ports: Dict[str, List[str]] = {}
 current_bridge: Optional[str] = None
 
-# Background asyncio loop used to run DB tasks
+# small bounded buffer for packets produced before shared_objects is ready
+_PACKET_BUFFER: List[Dict[str, Any]] = []
+_BUFFER_CAPACITY = 20000
+
+# Background asyncio loop used for internal tasks in this module (kept but not used for DB pool)
 async_loop = asyncio.new_event_loop()
 
 
@@ -60,11 +63,37 @@ threading.Thread(target=_start_async_loop, args=(async_loop,), daemon=True).star
 
 
 # -------------------------
-# Database init (pool)
+# Helpers for buffer draining
 # -------------------------
-db = DB(DB_DSN)
-# schedule pool init on background loop (fire-and-forget)
-asyncio.run_coroutine_threadsafe(db.init_pool(), async_loop)
+def drain_buffer_to_shared_db() -> None:
+    """
+    Attempt to schedule buffered packets for insertion on shared_objects.web_loop.
+    Call this from main.py after shared_objects.db and shared_objects.web_loop are initialized.
+    """
+    try:
+        web_loop = getattr(shared_objects, "web_loop", None)
+        web_db = getattr(shared_objects, "db", None)
+        if web_db is None or web_loop is None:
+            return
+
+        # schedule draining on the web loop to avoid blocking this thread
+        def _drain() -> None:
+            while _PACKET_BUFFER:
+                pkt = _PACKET_BUFFER.pop(0)
+                try:
+                    asyncio.run_coroutine_threadsafe(web_db.insert_packet(pkt), web_loop)
+                except Exception:
+                    # re-buffer first element and stop to avoid busy loop
+                    _PACKET_BUFFER.insert(0, pkt)
+                    break
+
+        try:
+            web_loop.call_soon_threadsafe(_drain)
+        except Exception:
+            # fallback: run directly (best-effort)
+            _drain()
+    except Exception:
+        logger.exception("Failed to drain packet buffer")
 
 
 # -------------------------
@@ -80,8 +109,6 @@ def _safe_get_attr(layer, attr: str):
 def parse_packet(pkt, bridge: str) -> None:
     """
     Parse a scapy Packet object into a normalized dict and schedule DB insert.
-    Keeps all information from the original implementation (fields, VLAN handling,
-    IP/IPv6/TCP/UDP/ARP handling and enum mapping).
     """
     pkt_iface = getattr(pkt, "sniffed_on", None)
     if not pkt_iface:
@@ -216,9 +243,18 @@ def parse_packet(pkt, bridge: str) -> None:
     if Raw in pkt and not pkt_info["protocol_name"]:
         pkt_info["protocol_name"] = "RAW"
 
-    # Submit DB insert to background loop (non-blocking)
+    # Submit DB insert to shared web loop if available, otherwise buffer
     try:
-        asyncio.run_coroutine_threadsafe(db.insert_packet(pkt_info), async_loop)
+        web_loop = getattr(shared_objects, "web_loop", None)
+        web_db = getattr(shared_objects, "db", None)
+        if web_db is not None and web_loop is not None:
+            asyncio.run_coroutine_threadsafe(web_db.insert_packet(pkt_info), web_loop)
+        else:
+            # buffer (bounded) until the web app initializes
+            _PACKET_BUFFER.append(pkt_info)
+            if len(_PACKET_BUFFER) > _BUFFER_CAPACITY:
+                # drop oldest packet
+                _PACKET_BUFFER.pop(0)
     except Exception:
         logger.exception("Failed to schedule DB insert")
 
@@ -419,9 +455,12 @@ def stop_afpacket_sniffer() -> None:
         fixed_bridge_ports.pop(current_bridge, None)
     current_bridge = None
 
-    # close db pool (schedule on async loop)
+    # close db pool if shared.web_loop is available; otherwise leave to main
     try:
-        asyncio.run_coroutine_threadsafe(db.close_pool(), async_loop)
+        web_loop = getattr(shared_objects, "web_loop", None)
+        web_db = getattr(shared_objects, "db", None)
+        if web_db is not None and web_loop is not None:
+            asyncio.run_coroutine_threadsafe(web_db.close_pool(), web_loop)
     except Exception:
         logger.exception("Failed to schedule DB pool close")