test backend typing packets
All checks were successful
Build and Deploy MITM Webserver / build (push) Successful in 11s

This commit is contained in:
2026-03-05 15:46:41 +01:00
parent 7812f8715f
commit d06c8adcdd
6 changed files with 104 additions and 41 deletions

View File

@@ -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...",
}
}

View File

@@ -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

View File

@@ -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 *
@@ -146,15 +148,32 @@ 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:

View File

@@ -334,7 +334,7 @@ export default function PacketViewer(): ReactElement {
),
},
{
title: 'Proto',
title: 'Protocol',
dataIndex: 'ip_proto',
key: 'ip_proto',
width: 100,

View File

@@ -18,7 +18,7 @@ export type IconProps = React.SVGProps<SVGSVGElement> & {
};
const FirewallIcon = forwardRef<SVGSVGElement, IconProps>(
({ size = 24, color = 'currentColor', title, ...rest }, ref) => {
({ size = 20, color = 'currentColor', title, ...rest }, ref) => {
// If user passed a numeric size, treat as px
const sizeValue = typeof size === 'number' ? `${size}px` : size;

View File

@@ -1,31 +1,9 @@
import { useEffect } from 'react';
import { fetchFullState } from '../api/apiClient';
import { useBackendAPI } from '../hooks/useBackendAPI';
import Title from 'antd/lib/typography/Title';
export default function Home() {
const { fullStateQuery } = useBackendAPI();
const { data, isLoading, isError, error } = fullStateQuery;
if (isLoading) return <div>Loading full state…</div>;
if (isError) return <div>Error: {(error as Error)?.message}</div>;
useEffect(() => {
fetchFullState().catch(() => {});
}, []);
return (
<div>
<h3>Full State</h3>
<div>
<strong>Interfaces:</strong> {data?.interfaces.length ?? 0}
</div>
<div>
<strong>Routes:</strong> {data?.routes.length ?? 0}
</div>
<div>
<strong>Bridges:</strong> {data?.bridges.length ?? 0}
</div>
<pre style={{ whiteSpace: 'pre-wrap' }}>{JSON.stringify(data, null, 2)}</pre>
<Title level={2}>TBD</Title>
</div>
);
}