Files
sdns-proxy/scripts/post_processing/add_extra_metrics_to_csv.py
T

259 lines
8.0 KiB
Python

#!/usr/bin/env python3
"""
Fast PCAP Preprocessor for DNS QoS Analysis.
Loads PCAP into memory, then uses binary search to match packets to query
windows. Direction is determined by a configured local IP (the netns veth IP).
"""
import argparse
import bisect
import csv
import shutil
import socket
import time
from pathlib import Path
from typing import Dict, List, NamedTuple
import dpkt
from dateutil import parser as date_parser
BANDWIDTH_COLUMNS = [
"bytes_sent",
"bytes_received",
"packets_sent",
"packets_received",
"total_bytes",
]
DEFAULT_LOCAL_IP = "192.168.100.2" # netns veth1 address
DEFAULT_PROVIDERS = ["adguard", "cloudflare", "google", "quad9"]
class Packet(NamedTuple):
timestamp: float
size: int
is_outbound: bool
def needs_processing(csv_path: Path) -> bool:
"""True if file lacks bandwidth columns OR all values are empty/zero."""
try:
with open(csv_path, "r", encoding="utf-8") as f:
reader = csv.DictReader(f)
if not reader.fieldnames:
return True
if not all(c in reader.fieldnames for c in BANDWIDTH_COLUMNS):
return True
for row in reader:
for col in BANDWIDTH_COLUMNS:
val = (row.get(col) or "").strip()
if val and val != "0":
return False # has real data
return True # columns exist but all empty/zero
except Exception:
return True
def parse_csv_timestamp(ts_str: str) -> float:
return date_parser.isoparse(ts_str).timestamp()
def load_pcap(pcap_path: Path, local_ip_bytes: bytes) -> List[Packet]:
"""Load PCAP into a list of Packets sorted by timestamp."""
print(" Loading PCAP...")
t0 = time.time()
packets: List[Packet] = []
with open(pcap_path, "rb") as f:
try:
reader = dpkt.pcap.Reader(f)
except ValueError:
f.seek(0)
reader = dpkt.pcapng.Reader(f)
for ts, buf in reader:
try:
eth = dpkt.ethernet.Ethernet(buf)
ip = eth.data
if not isinstance(ip, dpkt.ip.IP):
continue
if ip.src == local_ip_bytes:
is_outbound = True
elif ip.dst == local_ip_bytes:
is_outbound = False
else:
continue # not our traffic
packets.append(Packet(float(ts), len(buf), is_outbound))
except (dpkt.dpkt.NeedData, dpkt.dpkt.UnpackError, AttributeError):
continue
packets.sort(key=lambda p: p.timestamp)
print(f" Loaded {len(packets):,} packets in {time.time() - t0:.2f}s")
return packets
def load_csv_queries(csv_path: Path) -> List[Dict]:
queries = []
with open(csv_path, "r", encoding="utf-8") as f:
for row in csv.DictReader(f):
try:
start = parse_csv_timestamp(row["timestamp"])
duration = float(row["duration_ns"]) / 1e9
queries.append(
{"data": row, "start_time": start, "end_time": start + duration}
)
except Exception as e:
print(f" Warning: skipping row - {e}")
return queries
def match_packets(packets: List[Packet], queries: List[Dict]) -> int:
"""Assign bandwidth metrics to each query. Returns total matched packets."""
if not packets or not queries:
for q in queries:
q.update({c: 0 for c in BANDWIDTH_COLUMNS})
return 0
print(" Matching packets to queries...")
t0 = time.time()
timestamps = [p.timestamp for p in packets]
matched = 0
for q in queries:
lo = bisect.bisect_left(timestamps, q["start_time"])
hi = bisect.bisect_right(timestamps, q["end_time"])
bs = br = ps = pr = 0
for pkt in packets[lo:hi]:
if pkt.is_outbound:
bs += pkt.size
ps += 1
else:
br += pkt.size
pr += 1
q["bytes_sent"] = bs
q["bytes_received"] = br
q["packets_sent"] = ps
q["packets_received"] = pr
q["total_bytes"] = bs + br
matched += ps + pr
print(f" Matched {matched:,} packets in {time.time() - t0:.2f}s")
total_sent = sum(q["bytes_sent"] for q in queries)
total_recv = sum(q["bytes_received"] for q in queries)
with_data = sum(1 for q in queries if q["total_bytes"] > 0)
print(f" Total: {total_sent:,} B sent, {total_recv:,} B received")
print(f" Queries with data: {with_data}/{len(queries)}")
return matched
def write_csv(csv_path: Path, queries: List[Dict], backup: bool = True):
if backup and csv_path.exists():
bak = csv_path.with_suffix(".csv.bak")
if not bak.exists():
shutil.copy2(csv_path, bak)
print(f" Backup: {bak.name}")
original_fields = [
f for f in queries[0]["data"].keys() if f not in BANDWIDTH_COLUMNS
]
fieldnames = original_fields + BANDWIDTH_COLUMNS
with open(csv_path, "w", encoding="utf-8", newline="") as f:
writer = csv.DictWriter(f, fieldnames=fieldnames)
writer.writeheader()
for q in queries:
row = {k: q["data"][k] for k in original_fields}
for c in BANDWIDTH_COLUMNS:
row[c] = q[c]
writer.writerow(row)
print(f" Written: {csv_path.name}")
def process_provider(provider_path: Path, local_ip_bytes: bytes):
print(f"\n{'=' * 60}\nProcessing: {provider_path.name.upper()}\n{'=' * 60}")
processed = skipped = 0
total_time = 0.0
for csv_path in sorted(provider_path.glob("*.csv")):
name = csv_path.name.lower()
if ".bak" in name or name.endswith((".cpu.csv", ".mem.csv")):
continue
pcap_path = csv_path.with_suffix(".pcap")
if not pcap_path.exists():
print(f"\n{csv_path.name}: no matching PCAP")
continue
if not needs_processing(csv_path):
print(f"\n{csv_path.name}: already processed")
skipped += 1
continue
print(f"\n 📁 {csv_path.name}")
t0 = time.time()
packets = load_pcap(pcap_path, local_ip_bytes)
if not packets:
print(" ⚠ No usable packets in PCAP")
continue
queries = load_csv_queries(csv_path)
if not queries:
print(" ⚠ No valid queries in CSV")
continue
print(f" Loaded {len(queries):,} queries")
match_packets(packets, queries)
write_csv(csv_path, queries)
dt = time.time() - t0
total_time += dt
processed += 1
print(f" ✓ Completed in {dt:.2f}s")
print(
f"\n {provider_path.name}: {processed} processed, "
f"{skipped} skipped, {total_time:.2f}s"
)
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument(
"--local-ip",
default=DEFAULT_LOCAL_IP,
help=f"Local (netns veth) IP used to determine direction "
f"(default: {DEFAULT_LOCAL_IP})",
)
ap.add_argument("--results-dir", default="results", type=Path)
ap.add_argument("--providers", nargs="+", default=DEFAULT_PROVIDERS)
args = ap.parse_args()
local_ip_bytes = socket.inet_aton(args.local_ip)
print(f"\n{'=' * 60}\nDNS PCAP PREPROCESSOR\n{'=' * 60}")
print(f"Local IP: {args.local_ip}")
print(f"Results: {args.results_dir}")
if not args.results_dir.exists():
print(f"\n❌ Directory not found: {args.results_dir}")
return
t0 = time.time()
for provider in args.providers:
path = args.results_dir / provider
if path.exists():
process_provider(path, local_ip_bytes)
else:
print(f"\n⚠ Missing provider directory: {provider}")
total = time.time() - t0
print(f"\n{'=' * 60}")
print(f"✓ DONE in {total:.2f}s ({total / 60:.1f} min)")
print(f"{'=' * 60}\n")
if __name__ == "__main__":
main()