routeros-suricata-tzsp/app/forensics.py

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("<IHHIIII", 0xA1B2C3D4, 2, 4, 0, 0, 65535, 1)
_SAFE_LABEL_RE = re.compile(r"[^A-Za-z0-9_.-]+")
_VLAN_TYPES = {0x8100, 0x88A8, 0x9100}


class ForensicPcapRing:
    """Bounded pre-event packet buffer used to persist evidence only after blocks.

    Frames live in RAM for a short window. Nothing is written to persistent storage
    until ``capture_target`` is called after a successful RouterOS block action.
    """

    def __init__(
        self,
        mode: str,
        directory: str,
        *,
        window_seconds: int = 60,
        memory_mb: int = 64,
        max_files: int = 32,
        max_total_mb: int = 512,
    ) -> 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("<IIII", seconds, micros, length, len(frame)))
                    handle.write(frame[:length])
            os.replace(temp_path, target_path)
            os.chmod(target_path, 0o640)
        finally:
            try:
                temp_path.unlink(missing_ok=True)
            except OSError:
                pass

        self._prune_files()
        try:
            size = target_path.stat().st_size
        except OSError:
            size = 0
        return {
            "name": target_path.name,
            "path": str(target_path),
            "size_bytes": size,
            "packet_count": len(packets),
            "target": normalized,
        }

    def status(self) -> 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()