124 lines
4.6 KiB
Python
124 lines
4.6 KiB
Python
"""Prepare fixed NSL-KDD train/test tables for Module 6 notebooks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
RAW_DIR = HERE / "raw"
|
|
|
|
FEATURE_NAMES = [
|
|
"duration", "protocol_type", "service", "flag", "src_bytes", "dst_bytes",
|
|
"land", "wrong_fragment", "urgent", "hot", "num_failed_logins",
|
|
"logged_in", "num_compromised", "root_shell", "su_attempted", "num_root",
|
|
"num_file_creations", "num_shells", "num_access_files", "num_outbound_cmds",
|
|
"is_host_login", "is_guest_login", "count", "srv_count", "serror_rate",
|
|
"srv_serror_rate", "rerror_rate", "srv_rerror_rate", "same_srv_rate",
|
|
"diff_srv_rate", "srv_diff_host_rate", "dst_host_count",
|
|
"dst_host_srv_count", "dst_host_same_srv_rate", "dst_host_diff_srv_rate",
|
|
"dst_host_same_src_port_rate", "dst_host_srv_diff_host_rate",
|
|
"dst_host_serror_rate", "dst_host_srv_serror_rate", "dst_host_rerror_rate",
|
|
"dst_host_srv_rerror_rate",
|
|
]
|
|
|
|
CATEGORICAL_FEATURES = ["protocol_type", "service", "flag"]
|
|
|
|
ATTACK_GROUPS = {
|
|
"normal": "Normal",
|
|
"back": "DoS", "land": "DoS", "neptune": "DoS", "pod": "DoS",
|
|
"smurf": "DoS", "teardrop": "DoS", "apache2": "DoS", "mailbomb": "DoS",
|
|
"processtable": "DoS", "udpstorm": "DoS", "worm": "DoS",
|
|
"ipsweep": "Probe", "nmap": "Probe", "portsweep": "Probe", "satan": "Probe",
|
|
"mscan": "Probe", "saint": "Probe",
|
|
"ftp_write": "R2L", "guess_passwd": "R2L", "imap": "R2L",
|
|
"multihop": "R2L", "phf": "R2L", "spy": "R2L", "warezclient": "R2L",
|
|
"warezmaster": "R2L", "sendmail": "R2L", "named": "R2L",
|
|
"snmpgetattack": "R2L", "snmpguess": "R2L", "xlock": "R2L",
|
|
"xsnoop": "R2L", "httptunnel": "R2L",
|
|
"buffer_overflow": "U2R", "loadmodule": "U2R", "perl": "U2R",
|
|
"rootkit": "U2R", "ps": "U2R", "sqlattack": "U2R", "xterm": "U2R",
|
|
}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--full-train",
|
|
action="store_true",
|
|
help="Use KDDTrain+.txt instead of KDDTrain+_20Percent.txt.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def find_one(filename: str) -> Path:
|
|
matches = list(RAW_DIR.rglob(filename)) if RAW_DIR.exists() else []
|
|
if len(matches) != 1:
|
|
raise FileNotFoundError(
|
|
f"Expected exactly one {filename!r} below {RAW_DIR}. "
|
|
"Follow data/README.md and verify the extracted files."
|
|
)
|
|
return matches[0]
|
|
|
|
|
|
def load_nsl_kdd(path: Path) -> pd.DataFrame:
|
|
columns = FEATURE_NAMES + ["attack_label", "difficulty"]
|
|
frame = pd.read_csv(path, header=None, names=columns, low_memory=False)
|
|
frame["attack_label"] = (
|
|
frame["attack_label"].astype(str).str.strip().str.rstrip(".").str.lower()
|
|
)
|
|
|
|
unknown = sorted(set(frame["attack_label"]) - set(ATTACK_GROUPS))
|
|
if unknown:
|
|
raise ValueError(f"Unknown attack labels in {path.name}: {unknown}")
|
|
|
|
for column in CATEGORICAL_FEATURES:
|
|
frame[column] = frame[column].astype(str).str.strip().str.lower()
|
|
|
|
numeric_features = [name for name in FEATURE_NAMES if name not in CATEGORICAL_FEATURES]
|
|
for column in numeric_features + ["difficulty"]:
|
|
frame[column] = pd.to_numeric(frame[column], errors="raise")
|
|
|
|
frame["attack_group"] = frame["attack_label"].map(ATTACK_GROUPS)
|
|
frame["binary_label"] = (frame["attack_group"] != "Normal").astype(int)
|
|
return frame[FEATURE_NAMES + ["attack_label", "attack_group", "binary_label", "difficulty"]]
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
train_name = "KDDTrain+.txt" if args.full_train else "KDDTrain+_20Percent.txt"
|
|
train = load_nsl_kdd(find_one(train_name))
|
|
test = load_nsl_kdd(find_one("KDDTest+.txt"))
|
|
|
|
train_path = HERE / "nsl_kdd_train.csv"
|
|
test_path = HERE / "nsl_kdd_test.csv"
|
|
manifest_path = HERE / "feature_manifest.csv"
|
|
|
|
train.to_csv(train_path, index=False)
|
|
test.to_csv(test_path, index=False)
|
|
|
|
manifest = pd.DataFrame({
|
|
"feature": FEATURE_NAMES,
|
|
"feature_type": [
|
|
"categorical" if name in CATEGORICAL_FEATURES else "numeric"
|
|
for name in FEATURE_NAMES
|
|
],
|
|
"model_feature": True,
|
|
})
|
|
manifest.to_csv(manifest_path, index=False)
|
|
|
|
print(f"Created {train_path} with shape {train.shape}")
|
|
print(f"Created {test_path} with shape {test.shape}")
|
|
print(f"Created {manifest_path}")
|
|
print("Train attack groups:")
|
|
print(train["attack_group"].value_counts())
|
|
print("Test attack groups:")
|
|
print(test["attack_group"].value_counts())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|