netflower 1.0.0__py3-none-any.whl

netflower/__init__.py

@@ -0,0 +1,5 @@
+from ._convert import convert_pcap_to_csv
+from ._capture import capture_live
+
+__all__ = ["convert_pcap_to_csv", "capture_live"]
+__version__ = "1.0.0"

netflower/_capture.py

@@ -0,0 +1,222 @@
+"""
+Live packet capture using system libpcap via ctypes.
+
+Usage:
+    from netflower import capture_live
+
+    handle = capture_live("eth0", on_flow=lambda flow: print(flow), idle_timeout=30)
+    handle.start()
+    # ... do work ...
+    handle.stop()
+"""
+
+import ctypes
+import threading
+
+from ._libpcap import (
+    PcapHandler_cb, PktHdr,
+    pcap_open_live, pcap_loop, pcap_breakloop, pcap_close,
+    pcap_dump_open, pcap_dump, pcap_dump_close,
+)
+from ._parser import _parse_packet
+from ._session import FlowSession
+from ._constants import FLOW_TIMEOUT
+
+
+class CallbackWriter:
+    """Writer that calls on_flow(dict) instead of writing to a CSV file."""
+
+    def __init__(self, on_flow: callable) -> None:
+        self._on_flow = on_flow
+
+    def write(self, flow_dict: dict) -> None:
+        self._on_flow(flow_dict)
+
+
+class _PcapSavingWriter:
+    """
+    Wraps on_flow callable. When a flow is emitted, also dumps its raw packets
+    to a .pcap file in pcap_dir before clearing the buffer.
+    """
+
+    def __init__(self, on_flow: callable, pcap_dir: str, pcap_handle) -> None:
+        self._on_flow = on_flow
+        self._pcap_dir = pcap_dir
+        self._pcap_handle = pcap_handle
+        # flow_key -> list of (PktHdr, bytes)
+        self._buffers: dict[tuple, list] = {}
+
+    def write(self, flow_dict: dict) -> None:
+        key = (
+            flow_dict.get("protocol"),
+            flow_dict.get("src_ip"),
+            flow_dict.get("dst_ip"),
+            flow_dict.get("src_port"),
+            flow_dict.get("dst_port"),
+        )
+        packets = self._buffers.pop(key, [])
+        if packets:
+            ts_us = int(flow_dict.get("timestamp", 0) * 1_000_000)
+            fname = (
+                f"{self._pcap_dir}/"
+                f"{flow_dict.get('src_ip')}-{flow_dict.get('src_port')}-"
+                f"{flow_dict.get('dst_ip')}-{flow_dict.get('dst_port')}-"
+                f"{flow_dict.get('protocol')}-{ts_us}.pcap"
+            ).encode()
+            dumper = pcap_dump_open(self._pcap_handle, fname)
+            if dumper:
+                for hdr, data in packets:
+                    pcap_dump(dumper, ctypes.byref(hdr), data)
+                pcap_dump_close(dumper)
+        self._on_flow(flow_dict)
+
+    def buffer_packet(self, key: tuple, hdr: PktHdr, data: bytes) -> None:
+        if key not in self._buffers:
+            self._buffers[key] = []
+        self._buffers[key].append((hdr, data))
+
+
+class CaptureHandle:
+    """
+    Handle returned by capture_live(). Controls start/stop of capture thread.
+
+    Parameters
+    ----------
+    interface:    Network interface name as bytes (e.g. b"eth0").
+    on_flow:      Callable receiving a flow dict when a flow completes.
+    idle_timeout: Seconds of inactivity before a flow is emitted.
+    flow_timeout: Absolute max flow duration before forced emit.
+    save_pcap:    If True, save raw packets of each completed flow to pcap_dir.
+    pcap_dir:     Directory path for .pcap files. Required when save_pcap=True.
+    """
+
+    _GC_INTERVAL = 1000
+
+    def __init__(
+        self,
+        interface: bytes,
+        on_flow: callable,
+        idle_timeout: float,
+        flow_timeout: float,
+        save_pcap: bool,
+        pcap_dir: str | None,
+    ) -> None:
+        if save_pcap and pcap_dir is None:
+            raise ValueError("pcap_dir is required when save_pcap=True")
+        self._interface = interface
+        self._on_flow = on_flow
+        self._idle_timeout = idle_timeout
+        self._flow_timeout = flow_timeout
+        self._save_pcap = save_pcap
+        self._pcap_dir = pcap_dir
+        self._handle = None
+        self._thread: threading.Thread | None = None
+        self._session = None
+        self._writer = None
+        self._callback = None
+
+    def start(self) -> None:
+        errbuf = ctypes.create_string_buffer(256)
+        self._handle = pcap_open_live(self._interface, 65535, 1, 1000, errbuf)
+        if not self._handle:
+            raise RuntimeError(f"pcap_open_live failed: {errbuf.value.decode()}")
+
+        if self._save_pcap:
+            self._writer = _PcapSavingWriter(self._on_flow, self._pcap_dir, self._handle)
+        else:
+            self._writer = CallbackWriter(self._on_flow)
+
+        self._session = FlowSession(self._writer, self._flow_timeout)
+        self._pkt_counter = 0
+
+        @PcapHandler_cb
+        def _callback(user, pkthdr_p, pkt_data):
+            hdr = pkthdr_p.contents
+            ts = hdr.ts_sec + hdr.ts_usec / 1_000_000
+            buf = ctypes.string_at(pkt_data, hdr.caplen)
+            parsed = _parse_packet(buf)
+            if parsed is None:
+                return
+            src_ip, dst_ip, src_port, dst_port, proto, pkt_len, hdr_len, pay_len, flags, win = parsed
+            if self._save_pcap:
+                fwd_key = (proto, src_ip, dst_ip, src_port, dst_port)
+                bwd_key = (proto, dst_ip, src_ip, dst_port, src_port)
+                key = fwd_key if (fwd_key in self._session._flows
+                                  or bwd_key not in self._session._flows) else bwd_key
+                self._writer.buffer_packet(key, hdr, buf)
+            self._session.process(ts, src_ip, dst_ip, src_port, dst_port,
+                                  proto, pkt_len, hdr_len, pay_len, flags, win)
+            self._pkt_counter += 1
+            if self._pkt_counter % self._GC_INTERVAL == 0:
+                self._session.gc(ts)
+
+        # Keep reference to prevent ctypes GC
+        self._callback = _callback
+
+        self._thread = threading.Thread(
+            target=pcap_loop,
+            args=(self._handle, -1, self._callback, None),
+            daemon=True,
+        )
+        self._thread.start()
+
+    def stop(self) -> None:
+        if self._handle:
+            pcap_breakloop(self._handle)
+        if self._thread:
+            self._thread.join(timeout=5)
+        if self._session:
+            self._session.flush_all()
+        if self._handle:
+            pcap_close(self._handle)
+            self._handle = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        self.stop()
+
+
+def capture_live(
+    interface: str,
+    on_flow: callable,
+    idle_timeout: float = 30.0,
+    flow_timeout: float = FLOW_TIMEOUT,
+    save_pcap: bool = False,
+    pcap_dir: str | None = None,
+) -> CaptureHandle:
+    """
+    Capture live packets from a network interface and emit completed flows.
+
+    Requires root or CAP_NET_RAW privilege.
+
+    Parameters
+    ----------
+    interface:    Network interface name (e.g. "eth0").
+    on_flow:      Called with a flow dict (82 fields) when a flow completes.
+                  Same column names as convert_pcap_to_csv output.
+    idle_timeout: Seconds of inactivity before a flow is emitted (default 30).
+    flow_timeout: Absolute max flow duration in seconds (default 120).
+    save_pcap:    Save raw packets of each completed flow to a .pcap file.
+    pcap_dir:     Directory for .pcap files. Required when save_pcap=True.
+
+    Returns
+    -------
+    CaptureHandle — call .start() to begin capture, .stop() to end.
+
+    Example
+    -------
+    >>> handle = capture_live("eth0", on_flow=lambda f: print(f["flow_duration"]))
+    >>> handle.start()
+    >>> import time; time.sleep(60)
+    >>> handle.stop()
+    """
+    return CaptureHandle(
+        interface=interface.encode() if isinstance(interface, str) else interface,
+        on_flow=on_flow,
+        idle_timeout=idle_timeout,
+        flow_timeout=flow_timeout,
+        save_pcap=save_pcap,
+        pcap_dir=pcap_dir,
+    )

netflower/_constants.py

@@ -0,0 +1,18 @@
+FLOW_TIMEOUT = 120.0    # seconds of inactivity before a flow is flushed
+ACTIVE_TIMEOUT = 5.0    # gap > this separates active from idle periods
+CLUMP_TIMEOUT = 1.0     # gap > this starts a new subflow
+BULK_BOUND = 4          # minimum packets required to register a bulk transfer
+CSV_BUFFER_ROWS = 500   # rows buffered in memory before writing to disk
+
+# TCP flag bitmasks (RFC 793 / RFC 3168)
+TCP_FIN = 0x01
+TCP_SYN = 0x02
+TCP_RST = 0x04
+TCP_PSH = 0x08
+TCP_ACK = 0x10
+TCP_URG = 0x20
+TCP_ECE = 0x40
+TCP_CWR = 0x80
+
+FORWARD = 0
+BACKWARD = 1

netflower/_convert.py

@@ -0,0 +1,248 @@
+"""
+Main conversion pipeline: PCAP → FlowSession → CSV.
+
+Supports both pcap and pcapng formats.
+Packets are streamed one by one; no packet objects are retained after processing.
+
+Parallelism strategy (n_jobs > 1):
+  The main process reads packets sequentially and routes each one to a worker
+  subprocess based on a deterministic hash of the bidirectional flow key.
+  Each worker owns a disjoint subset of flows and writes to a temp CSV.
+  Temp CSVs are merged into the final output after all workers finish.
+"""
+
+import os
+import tempfile
+from multiprocessing import Process, Queue
+from typing import Iterator
+
+import dpkt
+import dpkt.pcapng
+
+from ._session import FlowSession
+from ._writer import CsvWriter
+from ._parser import _parse_packet
+from ._constants import FLOW_TIMEOUT, CSV_BUFFER_ROWS
+
+
+def _open_pcap(path: str):
+    """Return an iterable of (timestamp, raw_bytes) from a pcap or pcapng file."""
+    with open(path, "rb") as fh:
+        magic = fh.read(4)
+
+    if magic == b"\x0a\x0d\x0d\x0a":
+        return _iter_pcapng(path)
+    return _iter_pcap(path)
+
+
+def _iter_pcap(path: str) -> Iterator[tuple[float, bytes]]:
+    with open(path, "rb") as fh:
+        yield from dpkt.pcap.Reader(fh)
+
+
+def _iter_pcapng(path: str) -> Iterator[tuple[float, bytes]]:
+    with open(path, "rb") as fh:
+        yield from dpkt.pcapng.Reader(fh)
+
+
+def _route(src_ip: str, dst_ip: str, src_port: int, dst_port: int, proto: int, n_jobs: int) -> int:
+    """Return the worker index for a bidirectional flow (order-independent)."""
+    a = (src_ip, src_port)
+    b = (dst_ip, dst_port)
+    if a > b:
+        a, b = b, a
+    return hash((a, b, proto)) % n_jobs
+
+
+def _worker(task_queue: Queue, result_queue: Queue, output_path: str,
+            flow_timeout: float, buffer_rows: int) -> None:
+    """Subprocess: drain task_queue, track flows, write CSV to output_path."""
+    with CsvWriter(output_path, buffer_rows) as writer:
+        session = FlowSession(writer, flow_timeout)
+        while True:
+            item = task_queue.get()
+            if item is None:
+                break
+            kind = item[0]
+            if kind == "pkt":
+                _, ts, src_ip, dst_ip, src_port, dst_port, proto, pkt_len, hdr_len, pay_len, flags, win = item
+                session.process(ts, src_ip, dst_ip, src_port, dst_port, proto,
+                                pkt_len, hdr_len, pay_len, flags, win)
+            elif kind == "gc":
+                session.gc(item[1])
+        session.flush_all()
+    result_queue.put(writer.row_count)
+
+
+def _merge_csvs(sources: list[str], output_path: str) -> None:
+    import csv
+
+    header_written = False
+    fieldnames = None
+    with open(output_path, "w", newline="", buffering=1 << 16) as out_fh:
+        out_writer = None
+        for src in sources:
+            if not os.path.exists(src) or os.path.getsize(src) == 0:
+                continue
+            with open(src, "r", newline="") as in_fh:
+                reader = csv.DictReader(in_fh)
+                if not header_written:
+                    fieldnames = reader.fieldnames
+                    out_writer = csv.DictWriter(out_fh, fieldnames=fieldnames)
+                    out_writer.writeheader()
+                    header_written = True
+                for row in reader:
+                    out_writer.writerow(row)
+
+
+def convert_pcap_to_csv(
+    input_path: str,
+    output_path: str,
+    flow_timeout: float = FLOW_TIMEOUT,
+    gc_interval: int = 1000,
+    buffer_rows: int = CSV_BUFFER_ROWS,
+    n_jobs: int = 1,
+) -> int:
+    """
+    Convert a PCAP (or pcapng) file into a flow-based CSV file.
+
+    Each row in the output CSV represents one bidirectional network flow and
+    contains 82 features compatible with the CICFlowMeter feature set.
+
+    Parameters
+    ----------
+    input_path:   Path to the input .pcap or .pcapng file.
+    output_path:  Path for the output .csv file (created or overwritten).
+    flow_timeout: Seconds of inactivity before a flow is flushed (default 120).
+    gc_interval:  Run idle-flow garbage collection every N packets (default 1000).
+    buffer_rows:  Rows buffered in memory before a disk write (default 500).
+    n_jobs:       Number of worker processes (default 1).
+                  -1 uses all available CPUs.
+
+    Returns
+    -------
+    Number of flow rows written to the CSV.
+    """
+    if n_jobs == 0:
+        raise ValueError("n_jobs=0 is invalid. Use n_jobs=1 for single-process or n_jobs=-1 for all CPUs.")
+    if n_jobs == -1:
+        n_jobs = os.cpu_count() or 1
+    elif n_jobs < -1:
+        n_jobs = max(1, (os.cpu_count() or 1) + n_jobs + 1)
+
+    if n_jobs == 1:
+        return _convert_single(input_path, output_path, flow_timeout, gc_interval, buffer_rows)
+    return _convert_parallel(input_path, output_path, flow_timeout, gc_interval, buffer_rows, n_jobs)
+
+
+def _convert_single(
+    input_path: str,
+    output_path: str,
+    flow_timeout: float,
+    gc_interval: int,
+    buffer_rows: int,
+) -> int:
+    with CsvWriter(output_path, buffer_rows) as writer:
+        session = FlowSession(writer, flow_timeout)
+        pkt_counter = 0
+        last_ts = 0.0
+
+        for ts, buf in _open_pcap(input_path):
+            parsed = _parse_packet(buf)
+            if parsed is None:
+                continue
+
+            src_ip, dst_ip, src_port, dst_port, proto, pkt_len, hdr_len, pay_len, flags, win = parsed
+            session.process(ts, src_ip, dst_ip, src_port, dst_port, proto,
+                            pkt_len, hdr_len, pay_len, flags, win)
+
+            pkt_counter += 1
+            last_ts = ts
+
+            if pkt_counter % gc_interval == 0:
+                session.gc(last_ts)
+
+        session.flush_all()
+
+    return writer.row_count
+
+
+def _convert_parallel(
+    input_path: str,
+    output_path: str,
+    flow_timeout: float,
+    gc_interval: int,
+    buffer_rows: int,
+    n_jobs: int,
+) -> int:
+    tmp_dir = tempfile.mkdtemp(prefix="netflower_")
+    tmp_files = [os.path.join(tmp_dir, f"worker_{i}.csv") for i in range(n_jobs)]
+    task_queues = [Queue(maxsize=2000) for _ in range(n_jobs)]
+    result_queue: Queue = Queue()
+
+    workers = [
+        Process(
+            target=_worker,
+            args=(task_queues[i], result_queue, tmp_files[i], flow_timeout, buffer_rows),
+            daemon=True,
+        )
+        for i in range(n_jobs)
+    ]
+    for w in workers:
+        w.start()
+
+    try:
+        pkt_counter = 0
+        last_ts = 0.0
+
+        for ts, buf in _open_pcap(input_path):
+            parsed = _parse_packet(buf)
+            if parsed is None:
+                continue
+
+            src_ip, dst_ip, src_port, dst_port, proto, pkt_len, hdr_len, pay_len, flags, win = parsed
+            wid = _route(src_ip, dst_ip, src_port, dst_port, proto, n_jobs)
+            task_queues[wid].put((
+                "pkt", ts, src_ip, dst_ip, src_port, dst_port, proto,
+                pkt_len, hdr_len, pay_len, flags, win,
+            ))
+
+            pkt_counter += 1
+            last_ts = ts
+
+            if pkt_counter % gc_interval == 0:
+                for q in task_queues:
+                    q.put(("gc", last_ts))
+
+        for q in task_queues:
+            q.put(None)
+
+        for w in workers:
+            w.join()
+            if w.exitcode != 0:
+                raise RuntimeError(f"Worker process exited with code {w.exitcode}")
+
+    except Exception:
+        for q in task_queues:
+            try:
+                q.put_nowait(None)
+            except Exception:
+                pass
+        for w in workers:
+            w.terminate()
+        raise
+
+    total = sum(result_queue.get() for _ in range(n_jobs))
+    _merge_csvs(tmp_files, output_path)
+
+    for f in tmp_files:
+        try:
+            os.unlink(f)
+        except OSError:
+            pass
+    try:
+        os.rmdir(tmp_dir)
+    except OSError:
+        pass
+
+    return total