marcus/mitm-webserver
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

try to deduplicate db rows
All checks were successful
Build and Deploy MITM Webserver / traffic_target (push) Successful in 0s
Details
Build and Deploy MITM Webserver / build (push) Successful in 11s
Details

Browse Source
This commit is contained in:
malmert
2026-03-10 21:33:41 +01:00
parent 32aa17e0cf
commit c2fb35155c
2 changed files with 39 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
backend/src/utilities/database.py
View File

@@ -11,6 +11,8 @@ import asyncpg
from asyncpg.pool import Pool
from pydantic import ValidationError
from src.Models.etherType import EtherTypeEnum, ethertype_from_int
from src.Models.ip_protocol import protocol_from_number
from src.Models.packets import PacketDBModel
logger = logging.getLogger("af_packet_sniffer")
@@ -93,6 +95,24 @@ def _attach_derived_fields(payload: Dict[str, Any]) -> None:
if flow_id is not None:
payload["flow_id"] = flow_id
if payload.get("raw_present") is None:
payload["raw_present"] = payload.get("raw") is not None
if payload.get("eth_type") in (None, "") and payload.get("eth_type_raw") is not None:
try:
payload["eth_type"] = ethertype_from_int(int(payload["eth_type_raw"]))
except Exception:
payload["eth_type"] = EtherTypeEnum.UNKNOWN
if payload.get("ip_proto") in (None, "") and payload.get("ip_proto_raw") is not None:
try:
payload["ip_proto"] = protocol_from_number(int(payload["ip_proto_raw"]))
except Exception:
payload["ip_proto"] = int(payload["ip_proto_raw"])
if payload.get("app_master_protocol") in (None, "") and payload.get("app_protocol") not in (None, ""):
payload["app_master_protocol"] = payload.get("app_protocol")
class DatabasePool:
"""Asyncpg connection pool wrapper used by the packet APIs."""
@@ -178,19 +198,15 @@ class DatabasePool:
src_mac,
dst_mac,
eth_type_raw,
eth_type,
vlan_id,
src_ip,
dst_ip,
ip_proto_raw,
ip_proto,
src_port,
dst_port,
length,
raw_present,
capture_sources,
app_protocol,
app_master_protocol,
app_category,
app_confidence,
app_hostname,
@@ -202,8 +218,7 @@ class DatabasePool:
raw
) VALUES(
$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,
$21,$22,$23,$24,$25,$26,$27,$28,$29,$30,$31,$32,$33,$34,$35,$36,$37,$38::jsonb,
$39::jsonb,$40::jsonb,$41
$21,$22,$23,$24,$25,$26,$27,$28,$29,$30,$31,$32,$33::jsonb,$34::jsonb,$35::jsonb,$36
)
ON CONFLICT (correlation_key) DO UPDATE SET
updated_at = NOW(),
@@ -229,16 +244,13 @@ class DatabasePool:
src_mac = COALESCE(EXCLUDED.src_mac, packets.src_mac),
dst_mac = COALESCE(EXCLUDED.dst_mac, packets.dst_mac),
eth_type_raw = COALESCE(EXCLUDED.eth_type_raw, packets.eth_type_raw),
eth_type = COALESCE(EXCLUDED.eth_type, packets.eth_type),
vlan_id = COALESCE(EXCLUDED.vlan_id, packets.vlan_id),
src_ip = COALESCE(EXCLUDED.src_ip, packets.src_ip),
dst_ip = COALESCE(EXCLUDED.dst_ip, packets.dst_ip),
ip_proto_raw = COALESCE(EXCLUDED.ip_proto_raw, packets.ip_proto_raw),
ip_proto = COALESCE(EXCLUDED.ip_proto, packets.ip_proto),
src_port = COALESCE(EXCLUDED.src_port, packets.src_port),
dst_port = COALESCE(EXCLUDED.dst_port, packets.dst_port),
length = COALESCE(EXCLUDED.length, packets.length),
raw_present = COALESCE(EXCLUDED.raw_present, FALSE) OR COALESCE(packets.raw_present, FALSE),
capture_sources = (
SELECT ARRAY(
SELECT DISTINCT source
@@ -249,7 +261,6 @@ class DatabasePool:
)
),
app_protocol = COALESCE(EXCLUDED.app_protocol, packets.app_protocol),
app_master_protocol = COALESCE(EXCLUDED.app_master_protocol, packets.app_master_protocol),
app_category = COALESCE(EXCLUDED.app_category, packets.app_category),
app_confidence = COALESCE(EXCLUDED.app_confidence, packets.app_confidence),
app_hostname = COALESCE(EXCLUDED.app_hostname, packets.app_hostname),
@@ -280,19 +291,15 @@ class DatabasePool:
pkt_info.get("src_mac"),
pkt_info.get("dst_mac"),
pkt_info.get("eth_type_raw"),
_db_text(pkt_info.get("eth_type")),
pkt_info.get("vlan_id"),
pkt_info.get("src_ip"),
pkt_info.get("dst_ip"),
pkt_info.get("protocol_raw"),
_db_text(pkt_info.get("protocol_name") or pkt_info.get("protocol")),
pkt_info.get("src_port"),
pkt_info.get("dst_port"),
pkt_info.get("length"),
pkt_info.get("raw_present"),
pkt_info.get("capture_sources"),
pkt_info.get("app_protocol"),
pkt_info.get("app_master_protocol"),
pkt_info.get("app_category"),
pkt_info.get("app_confidence"),
pkt_info.get("app_hostname"),
@@ -349,22 +356,21 @@ class DatabasePool:
SET
updated_at = NOW(),
app_protocol = COALESCE(packets.app_protocol, $11),
app_master_protocol = COALESCE(packets.app_master_protocol, $12),
app_category = COALESCE(packets.app_category, $13),
app_confidence = COALESCE(packets.app_confidence, $14),
app_hostname = COALESCE(packets.app_hostname, $15),
app_is_encrypted = COALESCE(packets.app_is_encrypted, $16),
app_category = COALESCE(packets.app_category, $12),
app_confidence = COALESCE(packets.app_confidence, $13),
app_hostname = COALESCE(packets.app_hostname, $14),
app_is_encrypted = COALESCE(packets.app_is_encrypted, $15),
dpi_metadata = CASE
WHEN $17::jsonb IS NULL THEN packets.dpi_metadata
WHEN packets.dpi_metadata IS NULL THEN $17::jsonb
ELSE packets.dpi_metadata || $17::jsonb
WHEN $16::jsonb IS NULL THEN packets.dpi_metadata
WHEN packets.dpi_metadata IS NULL THEN $16::jsonb
ELSE packets.dpi_metadata || $16::jsonb
END,
capture_sources = (
SELECT ARRAY(
SELECT DISTINCT source
FROM unnest(
COALESCE(packets.capture_sources, ARRAY[]::text[]) ||
COALESCE($18::text[], ARRAY[]::text[])
COALESCE($17::text[], ARRAY[]::text[])
) AS source
)
)
@@ -383,13 +389,12 @@ class DatabasePool:
AND timestamp BETWEEN $9 AND $10
AND (
packets.app_protocol IS NULL
OR packets.app_master_protocol IS NULL
OR packets.app_category IS NULL
OR packets.app_confidence IS NULL
OR packets.app_hostname IS NULL
OR packets.app_is_encrypted IS NULL
OR ($17::jsonb IS NOT NULL)
OR (COALESCE(array_length($18::text[], 1), 0) > 0)
OR ($16::jsonb IS NOT NULL)
OR (COALESCE(array_length($17::text[], 1), 0) > 0)
)
RETURNING *
""",
@@ -404,7 +409,6 @@ class DatabasePool:
lower_bound,
upper_bound,
enrichment.get("app_protocol"),
enrichment.get("app_master_protocol"),
enrichment.get("app_category"),
enrichment.get("app_confidence"),
enrichment.get("app_hostname"),
@@ -465,29 +469,24 @@ class DatabasePool:
THEN COALESCE($7, packets.app_protocol)
ELSE packets.app_protocol
END,
app_master_protocol = CASE
WHEN packets.app_master_protocol IS NULL OR packets.app_master_protocol IN ('TCP', 'UDP', 'IP', 'IPv6', 'ETH')
THEN COALESCE($8, packets.app_master_protocol)
ELSE packets.app_master_protocol
END,
app_category = CASE
WHEN packets.app_category IS NULL OR packets.app_category IN ('Transport', 'Network', 'Protocol')
THEN COALESCE($9, packets.app_category)
THEN COALESCE($8, packets.app_category)
ELSE packets.app_category
END,
app_confidence = CASE
WHEN packets.app_protocol IS NULL OR packets.app_protocol IN ('TCP', 'UDP', 'IP', 'IPv6', 'ETH')
THEN COALESCE($10, packets.app_confidence)
THEN COALESCE($9, packets.app_confidence)
ELSE packets.app_confidence
END,
app_hostname = COALESCE(packets.app_hostname, $11),
app_is_encrypted = COALESCE(packets.app_is_encrypted, $12),
app_hostname = COALESCE(packets.app_hostname, $10),
app_is_encrypted = COALESCE(packets.app_is_encrypted, $11),
capture_sources = (
SELECT ARRAY(
SELECT DISTINCT source
FROM unnest(
COALESCE(packets.capture_sources, ARRAY[]::text[]) ||
COALESCE($13::text[], ARRAY[]::text[])
COALESCE($12::text[], ARRAY[]::text[])
) AS source
)
)
@@ -505,14 +504,12 @@ class DatabasePool:
AND (
packets.app_protocol IS NULL
OR packets.app_protocol IN ('TCP', 'UDP', 'IP', 'IPv6', 'ETH')
OR packets.app_master_protocol IS NULL
OR packets.app_master_protocol IN ('TCP', 'UDP', 'IP', 'IPv6', 'ETH')
OR packets.app_category IS NULL
OR packets.app_category IN ('Transport', 'Network', 'Protocol')
OR packets.app_confidence IS NULL
OR packets.app_hostname IS NULL
OR packets.app_is_encrypted IS NULL
OR (COALESCE(array_length($13::text[], 1), 0) > 0)
OR (COALESCE(array_length($12::text[], 1), 0) > 0)
)
RETURNING *
""",
@@ -523,7 +520,6 @@ class DatabasePool:
lower_bound,
upper_bound,
enrichment.get("app_protocol"),
enrichment.get("app_master_protocol"),
enrichment.get("app_category"),
enrichment.get("app_confidence"),
enrichment.get("app_hostname"),