from __future__ import annotations import ipaddress import os import re import struct import threading import time from collections import deque from datetime import datetime, timezone from pathlib import Path from typing import Any _PCAP_GLOBAL_HEADER = struct.pack(" None: self.mode = str(mode or "blocks").strip().lower() self.directory = Path(directory) self.window_seconds = max(1, int(window_seconds)) self.max_memory_bytes = max(1, int(memory_mb)) * 1024 * 1024 self.max_files = max(1, int(max_files)) self.max_total_bytes = max(1, int(max_total_mb)) * 1024 * 1024 self._frames: deque[tuple[float, bytes, frozenset[str]]] = deque() self._frame_bytes = 0 self._lock = threading.RLock() self.directory.mkdir(parents=True, exist_ok=True) @property def captures_blocks(self) -> bool: return self.mode == "blocks" def observe(self, frame: bytes) -> None: if not self.captures_blocks or not frame: return endpoints = _ethernet_ip_endpoints(frame) if not endpoints: return now = time.time() item = (now, bytes(frame), endpoints) with self._lock: self._frames.append(item) self._frame_bytes += len(item[1]) self._trim_locked(now) def capture_target(self, target: str, *, label: str = "block") -> dict[str, Any] | None: if not self.captures_blocks: return None try: normalized = str(ipaddress.ip_address(str(target).strip())) except ValueError: return None now = time.time() with self._lock: self._trim_locked(now) packets = [(timestamp, frame) for timestamp, frame, endpoints in self._frames if normalized in endpoints] if not packets: return None safe_label = _SAFE_LABEL_RE.sub("-", str(label or "block").strip()).strip("-._") or "block" safe_target = normalized.replace(":", "-") stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") target_path = self.directory / f"block-{stamp}-{safe_target}-{safe_label}.pcap" temp_path = target_path.with_suffix(target_path.suffix + ".tmp") try: with temp_path.open("wb") as handle: handle.write(_PCAP_GLOBAL_HEADER) for timestamp, frame in packets: seconds = int(timestamp) micros = int((timestamp - seconds) * 1_000_000) length = min(len(frame), 65535) handle.write(struct.pack(" dict[str, Any]: with self._lock: return { "mode": self.mode, "buffered_frames": len(self._frames), "buffered_bytes": self._frame_bytes, "window_seconds": self.window_seconds, "memory_bytes": self.max_memory_bytes, "max_files": self.max_files, "max_total_bytes": self.max_total_bytes, } def _trim_locked(self, now: float) -> None: cutoff = now - self.window_seconds while self._frames and (self._frames[0][0] < cutoff or self._frame_bytes > self.max_memory_bytes): _timestamp, frame, _endpoints = self._frames.popleft() self._frame_bytes -= len(frame) def _prune_files(self) -> None: try: files = sorted( (path for path in self.directory.glob("block-*.pcap") if path.is_file()), key=lambda path: path.stat().st_mtime, reverse=True, ) except OSError: return total = 0 for index, path in enumerate(files): try: size = path.stat().st_size except OSError: continue total += size if index >= self.max_files or total > self.max_total_bytes: try: path.unlink() except OSError: pass def _ethernet_ip_endpoints(frame: bytes) -> frozenset[str]: if len(frame) < 14: return frozenset() offset = 14 ether_type = int.from_bytes(frame[12:14], "big") while ether_type in _VLAN_TYPES: if len(frame) < offset + 4: return frozenset() ether_type = int.from_bytes(frame[offset + 2:offset + 4], "big") offset += 4 if ether_type == 0x0800: if len(frame) < offset + 20: return frozenset() version_ihl = frame[offset] if version_ihl >> 4 != 4 or (version_ihl & 0x0F) < 5: return frozenset() src = str(ipaddress.IPv4Address(frame[offset + 12:offset + 16])) dst = str(ipaddress.IPv4Address(frame[offset + 16:offset + 20])) return frozenset((src, dst)) if ether_type == 0x86DD: if len(frame) < offset + 40 or frame[offset] >> 4 != 6: return frozenset() src = str(ipaddress.IPv6Address(frame[offset + 8:offset + 24])) dst = str(ipaddress.IPv6Address(frame[offset + 24:offset + 40])) return frozenset((src, dst)) return frozenset()