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

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