262 lines
9.2 KiB
Python
262 lines
9.2 KiB
Python
"""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()
|