test backend typing packets

backend/src/Models/packets.py

@@ -0,0 +1,41 @@
+# src/models/packet.py
+from pydantic import BaseModel, Field
+from typing import Optional, Union
+
+class PacketDBModel(BaseModel):
+    id: Union[int, str]
+    timestamp: str = Field(..., description="ISO timestamp")
+    iface: str
+    src_mac: Optional[str] = None
+    dst_mac: Optional[str] = None
+    eth_type: Optional[Union[int, str]] = None
+    ip_proto: Optional[Union[int, str]] = None
+    src_ip: Optional[str] = None
+    dst_ip: Optional[str] = None
+    src_port: Optional[int] = None
+    dst_port: Optional[int] = None
+    vlan_id: Optional[int] = None
+    length: Optional[int] = None
+    raw_b64: Optional[str] = Field(None, description="Base64-encoded raw bytes")
+    direction: Optional[str] = None
+    packets: Optional[int] = None
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "id": 123,
+                "timestamp": "2026-03-05T12:34:56.789Z",
+                "iface": "eth0",
+                "src_mac": "aa:bb:cc:dd:ee:ff",
+                "dst_mac": "11:22:33:44:55:66",
+                "eth_type": "IPv4",
+                "ip_proto": "TCP",
+                "src_ip": "192.168.1.10",
+                "dst_ip": "192.168.1.1",
+                "src_port": 54321,
+                "dst_port": 80,
+                "vlan_id": None,
+                "length": 128,
+                "raw_b64": "BASE64...",
+            }
+        }

backend/src/network_sniffer.py

@@ -7,7 +7,7 @@ import socket
 import selectors
 import errno
 import struct
-from typing import Dict, List, Optional, Any
+from typing import Dict, List, Optional, Any, TypedDict, Union
 from uuid import uuid4
 
 # NOTE: ensure this path points to your shared runtime module
@@ -50,8 +50,33 @@ sessions: Dict[str, Dict[str, Any]] = {}
 #   "ports": List[str],
 # }
 
+# -------------------------
+# PacketInfo typing
+# -------------------------
+class PacketInfo(TypedDict, total=False):
+    """
+    TypedDict for the parsed packet info produced by parse_packet.
+    Fields marked optional (total=False) for flexibility across contexts.
+    """
+    iface: str
+    length: int
+    raw: bytes  # original raw bytes (kept for buffering; DB helper may convert to base64)
+    src_mac: Optional[str]
+    dst_mac: Optional[str]
+    eth_type_raw: Optional[int]
+    eth_type: Optional[Union[str, EtherTypeEnum]]
+    vlan_id: Optional[int]
+    protocol_raw: Optional[int]
+    protocol: Optional[Union[str, IPProtocolEnum, int]]
+    protocol_name: Optional[str]
+    src_ip: Optional[str]
+    dst_ip: Optional[str]
+    src_port: Optional[int]
+    dst_port: Optional[int]
+
+
 # small bounded buffer for packets produced before shared_objects is ready
-_PACKET_BUFFER: List[Dict[str, Any]] = []
+_PACKET_BUFFER: List[PacketInfo] = []
 _BUFFER_CAPACITY = 20000
 
 # Background asyncio loop used for internal tasks in this module (kept but not used for DB pool)
@@ -113,7 +138,7 @@ def _safe_get_attr(layer, attr: str):
 
 def parse_packet(pkt, bridge_label: str) -> None:
     """
-    Parse a scapy Packet object into a normalized dict and schedule DB insert.
+    Parse a scapy Packet object into a normalized PacketInfo and schedule DB insert.
     bridge_label indicates whether the packet was captured as part of a bridge-snapshot or single-interface.
     """
     pkt_iface = getattr(pkt, "sniffed_on", None)
@@ -122,7 +147,7 @@ def parse_packet(pkt, bridge_label: str) -> None:
 
     logger.debug("Packet captured on %s (bridge_label %s)", pkt_iface, bridge_label)
 
-    pkt_info: Dict[str, Any] = {
+    pkt_info: PacketInfo = {
         "iface": pkt_iface,
         "length": len(pkt),
         "raw": bytes(pkt),
@@ -210,7 +235,7 @@ def parse_packet(pkt, bridge_label: str) -> None:
         elif proto_num == 1 and ICMP in pkt:
             pkt_info["protocol_name"] = "ICMP"
         else:
-            if pkt_info["protocol_name"] is None:
+            if pkt_info.get("protocol_name") is None:
                 pkt_info["protocol_name"] = f"IP_PROTO_{proto_num}" if proto_num is not None else None
 
     # IPv6
@@ -242,11 +267,11 @@ def parse_packet(pkt, bridge_label: str) -> None:
         elif ICMPv6Unknown in pkt:
             pkt_info["protocol_name"] = "ICMPv6"
         else:
-            if pkt_info["protocol_name"] is None:
+            if pkt_info.get("protocol_name") is None:
                 pkt_info["protocol_name"] = f"IPV6_PROTO_{nh}" if nh is not None else None
 
     # Raw fallback label
-    if Raw in pkt and not pkt_info["protocol_name"]:
+    if Raw in pkt and not pkt_info.get("protocol_name"):
         pkt_info["protocol_name"] = "RAW"
 
     # Submit DB insert to shared web loop if available, otherwise buffer

backend/src/utilities/database.py

@@ -3,10 +3,14 @@ import logging
 import base64
 import asyncio
 from typing import Dict, List, Optional, Any
+from pydantic import ValidationError
+
 
 import asyncpg
 from asyncpg.pool import Pool
 
+from src.Models.packets import PacketDBModel
+
 # ---- Logging ----------------------------------------------------------
 logger = logging.getLogger("af_packet_sniffer")
 
@@ -124,18 +128,16 @@ class DatabasePool:
             except Exception:
                 logger.exception("Failed to publish pkt_info to broadcaster")
 
-    async def fetch_latest(self, limit: int) -> List[Dict[str, Any]]:
+    async def fetch_latest(self, limit: int) -> List[PacketDBModel]:
         """
         Fetch the latest `limit` packets (newest first).
 
-        Returns a list of dicts. If the `raw` column is binary it is converted
-        to `raw_b64` (base64 string) and `raw` is removed.
+        Returns a list of PacketDBModel. Converts raw bytes -> raw_b64 for JSON-safe output.
         """
         if self._pool is None:
             await self.init_pool()
 
         async with self._pool.acquire() as conn:
-            # Select explicit columns to ensure predictable dict keys
             rows = await conn.fetch(
                 """
                 SELECT *
@@ -145,16 +147,33 @@ class DatabasePool:
                 """,
                 limit,
             )
-            
-        out: List[Dict[str, Any]] = []
+
+        out: List[PacketDBModel] = []
+
         for r in rows:
             d = dict(r)
+
+            # convert byte raw -> base64 string (and remove raw)
             raw_val = d.get("raw")
             if isinstance(raw_val, (bytes, bytearray)):
                 d["raw_b64"] = base64.b64encode(raw_val).decode("ascii")
                 d.pop("raw", None)
-            out.append(d)
-            logger.exception(d)
+
+
+            # Validate/construct Pydantic model
+            try:
+                packet_model = PacketDBModel(**d)
+            except ValidationError as ve:
+                # Log and skip invalid rows (or handle otherwise)
+                logger.warning(
+                    "Skipping DB row that failed PacketDBModel validation (id=%s): %s",
+                    d.get("id"),
+                    ve,
+                )
+                continue
+
+            out.append(packet_model)
+
         return out
     
     async def clear_all_packets(self, reset_identity: bool = True) -> bool: