"""Create a bounded, classroom-ready flow table from one CTU-13 .binetflow file.""" from __future__ import annotations import argparse from pathlib import Path import numpy as np import pandas as pd HERE = Path(__file__).resolve().parent RAW_DIR = HERE / "raw" REQUIRED_RAW_COLUMNS = [ "StartTime", "Dur", "Proto", "SrcAddr", "Sport", "DstAddr", "Dport", "State", "sTos", "dTos", "TotPkts", "TotBytes", "SrcBytes", "Label", ] # CTU-13 is distributed in more than one Argus text schema. The compact # 15-column export contains ``Dir``. The extended export supplied with some # scenarios omits it but provides per-direction packet and byte counters. OPTIONAL_RAW_COLUMNS = ["Dir", "SrcPkts", "DstPkts", "DstBytes"] RENAME_MAP = { "StartTime": "start_time", "Dur": "duration", "Proto": "protocol", "SrcAddr": "src_addr", "Sport": "src_port", "Dir": "direction", "DstAddr": "dst_addr", "Dport": "dst_port", "State": "state", "sTos": "src_tos", "dTos": "dst_tos", "TotPkts": "total_packets", "TotBytes": "total_bytes", "SrcBytes": "src_bytes", "DstBytes": "dst_bytes", "Label": "raw_label", } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "input", nargs="?", type=Path, help="Path to one labeled CTU-13 .binetflow file. If omitted, data/raw is searched.", ) parser.add_argument( "--per-class", type=int, default=20_000, help="Maximum sampled rows for each of Botnet, Normal, and Background.", ) parser.add_argument( "--chunk-size", type=int, default=200_000, help="Rows read at a time; reduce this value on low-memory computers.", ) parser.add_argument( "--output", type=Path, default=HERE / "ctu13_botnet_flows.csv", help="Prepared CSV destination.", ) return parser.parse_args() def find_input(requested: Path | None) -> Path: if requested is not None: path = requested.expanduser().resolve() if not path.is_file(): raise FileNotFoundError(f"CTU-13 input file not found: {path}") return path matches = sorted(RAW_DIR.rglob("*.binetflow")) if RAW_DIR.exists() else [] if len(matches) != 1: raise FileNotFoundError( f"Expected exactly one .binetflow file below {RAW_DIR}, found {len(matches)}. " "Pass the input path explicitly when more than one file is present." ) return matches[0] def label_family(labels: pd.Series) -> pd.Series: text = labels.fillna("").astype(str) return pd.Series( np.select( [ text.str.contains("Botnet", case=False, regex=False), text.str.contains("Normal", case=False, regex=False), text.str.contains("Background", case=False, regex=False), ], ["Botnet", "Normal", "Background"], default="Other", ), index=labels.index, dtype="string", ) def parse_port(values: pd.Series) -> pd.Series: text = values.fillna("").astype(str).str.strip().str.lower() decimal = pd.to_numeric(text.where(~text.str.startswith("0x")), errors="coerce") hexadecimal = text.where(text.str.startswith("0x")).map( lambda value: int(value, 16) if isinstance(value, str) and value else np.nan ) return decimal.fillna(hexadecimal).astype("Float64") def inspect_schema(path: Path) -> tuple[list[str], str]: """Validate a compact or extended CTU-13 header and select needed fields.""" header = pd.read_csv(path, nrows=0).columns.tolist() missing = sorted(set(REQUIRED_RAW_COLUMNS) - set(header)) if missing: raise ValueError( f"Missing required CTU-13 columns: {missing}. Available columns: {header}" ) selected = REQUIRED_RAW_COLUMNS + [ column for column in OPTIONAL_RAW_COLUMNS if column in header ] if "Dir" in header: direction_mode = "source Dir column" elif {"SrcPkts", "DstPkts"}.issubset(header): direction_mode = "derived from SrcPkts/DstPkts" else: direction_mode = "constant unknown (no Dir or packet-direction counters)" print(f"Detected {len(header)}-column CTU-13 schema; direction is {direction_mode}.") return selected, direction_mode def ensure_direction(frame: pd.DataFrame) -> pd.DataFrame: """Return a frame with a categorical Dir field for either CTU-13 schema.""" if "Dir" in frame.columns: frame["Dir"] = frame["Dir"].fillna("unknown").astype(str).str.strip() frame.loc[frame["Dir"].eq(""), "Dir"] = "unknown" return frame if {"SrcPkts", "DstPkts"}.issubset(frame.columns): src_packets = pd.to_numeric(frame["SrcPkts"], errors="coerce").fillna(0) dst_packets = pd.to_numeric(frame["DstPkts"], errors="coerce").fillna(0) frame["Dir"] = np.select( [ (src_packets > 0) & (dst_packets > 0), (src_packets > 0) & (dst_packets <= 0), (src_packets <= 0) & (dst_packets > 0), ], ["bidirectional", "source_only", "destination_only"], default="unknown", ) else: frame["Dir"] = "unknown" return frame def bounded_sample(path: Path, per_class: int, chunk_size: int) -> pd.DataFrame: if per_class < 1 or chunk_size < 1: raise ValueError("--per-class and --chunk-size must be positive integers.") kept = {family: pd.DataFrame() for family in ("Botnet", "Normal", "Background")} rng = np.random.default_rng(42) saw_header = False selected_columns, _ = inspect_schema(path) for chunk_number, chunk in enumerate( pd.read_csv( path, usecols=selected_columns, chunksize=chunk_size, low_memory=False, ), start=1, ): saw_header = True chunk = ensure_direction(chunk) output_columns = REQUIRED_RAW_COLUMNS + ["Dir"] if "DstBytes" in chunk.columns: output_columns.append("DstBytes") chunk = chunk[output_columns].copy() chunk["label_family"] = label_family(chunk["Label"]) for family in kept: part = chunk.loc[chunk["label_family"].eq(family)].copy() if part.empty: continue part["_sample_key"] = rng.random(len(part)) candidates = pd.concat([kept[family], part], ignore_index=True) kept[family] = candidates.nsmallest(per_class, "_sample_key") counts = {family: len(frame) for family, frame in kept.items()} print(f"Processed chunk {chunk_number}: retained {counts}") if not saw_header: raise ValueError(f"No rows found in {path}") sampled = pd.concat(kept.values(), ignore_index=True).drop(columns="_sample_key") if sampled.empty: raise ValueError("No Botnet, Normal, or Background labels were found.") return sampled def engineer_features(sampled: pd.DataFrame) -> pd.DataFrame: frame = sampled.rename(columns=RENAME_MAP).copy() frame["start_time"] = pd.to_datetime(frame["start_time"], errors="coerce") numeric = [ "duration", "src_tos", "dst_tos", "total_packets", "total_bytes", "src_bytes" ] if "dst_bytes" in frame.columns: numeric.append("dst_bytes") for column in numeric: frame[column] = pd.to_numeric(frame[column], errors="coerce") frame["src_port"] = parse_port(frame["src_port"]) frame["dst_port"] = parse_port(frame["dst_port"]) derived_dst_bytes = (frame["total_bytes"] - frame["src_bytes"]).clip(lower=0) if "dst_bytes" in frame.columns: frame["dst_bytes"] = frame["dst_bytes"].fillna(derived_dst_bytes).clip(lower=0) else: frame["dst_bytes"] = derived_dst_bytes frame["bytes_per_packet"] = frame["total_bytes"] / frame["total_packets"].clip(lower=1) frame["src_byte_ratio"] = frame["src_bytes"] / frame["total_bytes"].clip(lower=1) frame["start_hour"] = frame["start_time"].dt.hour frame["is_well_known_dst_port"] = frame["dst_port"].between(0, 1023).astype("Int64") frame["binary_label"] = frame["label_family"].map( {"Normal": 0, "Botnet": 1, "Background": -1} ).astype(int) ordered = [ "start_time", "duration", "protocol", "src_addr", "src_port", "direction", "dst_addr", "dst_port", "state", "src_tos", "dst_tos", "total_packets", "total_bytes", "src_bytes", "dst_bytes", "bytes_per_packet", "src_byte_ratio", "start_hour", "is_well_known_dst_port", "raw_label", "label_family", "binary_label", ] return frame[ordered].sort_values("start_time", na_position="last").reset_index(drop=True) def main() -> None: args = parse_args() input_path = find_input(args.input) output_path = args.output.expanduser().resolve() output_path.parent.mkdir(parents=True, exist_ok=True) sampled = bounded_sample(input_path, args.per_class, args.chunk_size) prepared = engineer_features(sampled) prepared.to_csv(output_path, index=False) print(f"Created {output_path} with shape {prepared.shape}") print(prepared["label_family"].value_counts()) print("Background is retained for triage but excluded from supervised ground-truth metrics.") if __name__ == "__main__": main()