scripts/prepare_dataset.py

import argparse
import hashlib
import json
import os
import struct
import sys
import urllib.request
import zipfile
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from app.engine.ml_classifier import extract_features, extract_features_from_ember_record, FEATURE_NAMES

DATASET_DIR = Path(__file__).resolve().parent.parent / "datasets"
RAW_DIR = DATASET_DIR / "raw"
PROCESSED_DIR = DATASET_DIR / "processed"

EMBER_URL = "https://ember.elastic.co/ember_dataset_2024.tar.bz2"
BODMAS_INFO_URL = "https://whyisyoung.github.io/BODMAS/"

BENIGN_SCAN_DIRS = [
    Path(r"C:\Windows\System32"),
    Path(r"C:\Program Files"),
    Path(r"C:\Program Files (x86)"),
]

PE_EXTENSIONS = {".exe", ".dll", ".sys", ".ocx", ".scr", ".cpl", ".drv"}


def ensure_dirs():
    for d in [RAW_DIR, PROCESSED_DIR, RAW_DIR / "benign", RAW_DIR / "malicious"]:
        d.mkdir(parents=True, exist_ok=True)


def is_pe_file(filepath: Path) -> bool:
    try:
        with open(filepath, "rb") as f:
            magic = f.read(2)
            return magic == b"MZ"
    except (PermissionError, OSError):
        return False


def collect_benign_samples(max_samples: int = 50000) -> list[Path]:
    print(f"\n[1/4] Collecting benign PE samples from system directories...")
    collected = []
    seen_hashes = set()

    for scan_dir in BENIGN_SCAN_DIRS:
        if not scan_dir.exists():
            print(f"  Skipping {scan_dir} (not found)")
            continue

        print(f"  Scanning {scan_dir}...")
        try:
            for root, dirs, files in os.walk(scan_dir):
                dirs[:] = [d for d in dirs if d not in {"WinSxS", "Temp", "temp"}]
                for fname in files:
                    if len(collected) >= max_samples:
                        break
                    fpath = Path(root) / fname
                    if fpath.suffix.lower() not in PE_EXTENSIONS:
                        continue
                    if not is_pe_file(fpath):
                        continue
                    try:
                        file_hash = hashlib.sha256(fpath.read_bytes()).hexdigest()
                        if file_hash in seen_hashes:
                            continue
                        seen_hashes.add(file_hash)
                        collected.append(fpath)
                    except (PermissionError, OSError):
                        continue
                if len(collected) >= max_samples:
                    break
        except (PermissionError, OSError):
            continue

    print(f"  Collected {len(collected)} unique benign PE files")
    return collected


def download_ember(force: bool = False) -> Path | None:
    ember_dir = RAW_DIR / "ember2024"
    if ember_dir.exists() and not force and any(ember_dir.iterdir()):
        print(f"\n[2/4] EMBER2024 already present at {ember_dir}")
        return ember_dir

    print(f"\n[2/4] EMBER2024 Dataset Download")
    print(f"  The EMBER2024 dataset (~2GB) must be downloaded from:")
    print(f"  {EMBER_URL}")
    print(f"  ")
    print(f"  To download automatically, run:")
    print(f"    python scripts/prepare_dataset.py --download-ember")
    print(f"  ")
    print(f"  Or manually download and extract to: {ember_dir}")

    ember_dir.mkdir(parents=True, exist_ok=True)

    if force:
        tar_path = RAW_DIR / "ember_dataset_2024.tar.bz2"
        print(f"  Downloading EMBER2024 to {tar_path}...")
        try:
            urllib.request.urlretrieve(EMBER_URL, tar_path)
            print(f"  Download complete. Extracting...")
            import tarfile
            with tarfile.open(tar_path, "r:bz2") as tar:
                tar.extractall(path=ember_dir)
            print(f"  Extraction complete.")
            return ember_dir
        except Exception as e:
            print(f"  Download failed: {e}")
            print(f"  Please download manually from {EMBER_URL}")
            return None

    return None


def load_ember2024_from_zip(
    zip_path: Path,
    split: str = "train",
    max_per_week: int = 5000,
    val_weeks: int = 8,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None:
    if not zip_path.exists():
        print(f"  EMBER2024 zip not found: {zip_path}")
        return None

    prefix = f"Win64_{split}/"
    print(f"  Loading EMBER2024 from {zip_path.name} (split={split}, max_per_week={max_per_week})")

    with zipfile.ZipFile(zip_path) as zf:
        week_files = sorted(n for n in zf.namelist() if n.startswith(prefix) and n.endswith(".jsonl"))
        print(f"  Found {len(week_files)} weekly files")

        if split == "train" and val_weeks > 0:
            train_files = week_files[:-val_weeks]
            val_files = week_files[-val_weeks:]
        else:
            train_files = week_files
            val_files = []

        def _load_weeks(files: list[str], label_str: str) -> tuple[list, list]:
            X_list, y_list = [], []
            for i, fname in enumerate(files):
                week_X, week_y = [], []
                with zf.open(fname) as f:
                    for line in f:
                        record = json.loads(line)
                        lbl = record.get("label", -1)
                        if lbl == -1:
                            continue
                        feat = extract_features_from_ember_record(record)
                        if feat is not None:
                            week_X.append(feat)
                            week_y.append(lbl)
                        if len(week_X) >= max_per_week:
                            break

                X_list.extend(week_X)
                y_list.extend(week_y)
                benign = week_y.count(0)
                mal = week_y.count(1)
                print(f"    {label_str} week {i+1}/{len(files)}: {len(week_y)} samples (b={benign}, m={mal})")

            return X_list, y_list

        print(f"  Loading {len(train_files)} train weeks...")
        X_tr, y_tr = _load_weeks(train_files, "train")

        X_val, y_val = [], []
        if val_files:
            print(f"  Loading {len(val_files)} val weeks...")
            X_val, y_val = _load_weeks(val_files, "val")

    if not X_tr:
        print("  No training samples extracted.")
        return None

    X_train = np.array(X_tr, dtype=np.float64)
    y_train = np.array(y_tr, dtype=np.int64)
    X_val_arr = np.array(X_val, dtype=np.float64) if X_val else np.empty((0, len(FEATURE_NAMES)))
    y_val_arr = np.array(y_val, dtype=np.int64) if y_val else np.empty(0, dtype=np.int64)

    b_tr = int(np.sum(y_train == 0))
    m_tr = int(np.sum(y_train == 1))
    print(f"  Train: {len(y_train)} total (benign={b_tr}, malicious={m_tr})")
    if len(y_val_arr):
        b_v = int(np.sum(y_val_arr == 0))
        m_v = int(np.sum(y_val_arr == 1))
        print(f"  Val:   {len(y_val_arr)} total (benign={b_v}, malicious={m_v})")

    return X_train, y_train, X_val_arr, y_val_arr


def extract_pe_features(file_paths: list[Path], label: int) -> tuple[list[np.ndarray], list[int]]:
    X_list = []
    y_list = []
    skipped = 0

    for fpath in file_paths:
        try:
            file_bytes = fpath.read_bytes()
            features = extract_features(file_bytes)
            if features is not None:
                X_list.append(features)
                y_list.append(label)
            else:
                skipped += 1
        except (PermissionError, OSError, Exception):
            skipped += 1

    print(f"  Extracted features from {len(X_list)} files ({skipped} skipped)")
    return X_list, y_list


def generate_synthetic_vibeware(count: int = 5000, seed: int = 42) -> tuple[np.ndarray, np.ndarray]:
    print(f"\n[3/4] Generating {count} synthetic vibeware feature vectors...")
    rng = np.random.RandomState(seed)

    X_list = []
    for _ in range(count):
        variant = rng.choice(["nim", "zig", "rust", "go"])

        if variant == "nim":
            sections = rng.choice([3, 4, 5])
            max_entropy = rng.uniform(6.8, 7.6)
            mean_entropy = rng.uniform(5.2, 6.8)
            import_count = rng.randint(15, 80)
            dll_count = rng.randint(3, 12)
        elif variant == "zig":
            sections = rng.choice([2, 3, 4])
            max_entropy = rng.uniform(6.5, 7.4)
            mean_entropy = rng.uniform(5.0, 6.5)
            import_count = rng.randint(10, 50)
            dll_count = rng.randint(2, 8)
        elif variant == "rust":
            sections = rng.choice([4, 5, 6, 7])
            max_entropy = rng.uniform(6.2, 7.2)
            mean_entropy = rng.uniform(4.8, 6.2)
            import_count = rng.randint(30, 150)
            dll_count = rng.randint(5, 15)
        else:
            sections = rng.choice([3, 4, 5, 6])
            max_entropy = rng.uniform(6.0, 7.0)
            mean_entropy = rng.uniform(4.5, 6.0)
            import_count = rng.randint(20, 100)
            dll_count = rng.randint(4, 10)

        features = np.array([
            sections,
            rng.choice([224, 240]),
            rng.choice([0x0000, 0x0040, 0x0020, 0x8000]),
            rng.randint(1000, 50000),
            rng.randint(50000, 1500000),
            rng.randint(0, 80000),
            rng.randint(0x1000, 0x80000),
            rng.choice([0x00400000, 0x10000000, 0x140000000]),
            max_entropy,
            mean_entropy,
            rng.randint(0, 1024),
            rng.randint(50000, 3000000),
            import_count,
            dll_count,
            rng.choice([0, 1], p=[0.7, 0.3]),
        ], dtype=np.float64)

        X_list.append(features)

    X = np.array(X_list)
    y = np.ones(count, dtype=np.int64)
    print(f"  Generated {count} synthetic vibeware samples (Nim/Zig/Rust/Go profiles)")
    return X, y


def split_and_save(
    X: np.ndarray,
    y: np.ndarray,
    output_dir: Path,
    train_ratio: float = 0.7,
    val_ratio: float = 0.15,
    test_ratio: float = 0.15,
    seed: int = 42,
):
    print(f"\n[4/4] Splitting dataset ({len(X)} samples)...")
    print(f"  Ratios: train={train_ratio}, val={val_ratio}, test={test_ratio}")

    X_train, X_temp, y_train, y_temp = train_test_split(
        X, y, test_size=(val_ratio + test_ratio), random_state=seed, stratify=y
    )
    relative_test = test_ratio / (val_ratio + test_ratio)
    X_val, X_test, y_val, y_test = train_test_split(
        X_temp, y_temp, test_size=relative_test, random_state=seed, stratify=y_temp
    )

    output_dir.mkdir(parents=True, exist_ok=True)

    for name, X_split, y_split in [("train", X_train, y_train), ("val", X_val, y_val), ("test", X_test, y_test)]:
        np.save(output_dir / f"X_{name}.npy", X_split)
        np.save(output_dir / f"y_{name}.npy", y_split)

        df = pd.DataFrame(X_split, columns=FEATURE_NAMES)
        df["label"] = y_split
        df.to_parquet(output_dir / f"{name}.parquet", index=False)
        df.to_csv(output_dir / f"{name}.csv", index=False)

        benign_count = int(np.sum(y_split == 0))
        malicious_count = int(np.sum(y_split == 1))
        print(f"  {name:5s}: {len(y_split):6d} samples (benign={benign_count}, malicious={malicious_count})")

    manifest = {
        "total_samples": int(len(y)),
        "train_samples": int(len(y_train)),
        "val_samples": int(len(y_val)),
        "test_samples": int(len(y_test)),
        "feature_names": FEATURE_NAMES,
        "num_features": len(FEATURE_NAMES),
        "class_distribution": {
            "benign": int(np.sum(y == 0)),
            "malicious": int(np.sum(y == 1)),
        },
        "splits": {
            "train_ratio": train_ratio,
            "val_ratio": val_ratio,
            "test_ratio": test_ratio,
        },
    }
    with open(output_dir / "manifest.json", "w") as f:
        json.dump(manifest, f, indent=2)

    print(f"\n  Dataset saved to {output_dir}")
    print(f"  Files: X_{{train,val,test}}.npy, y_{{train,val,test}}.npy")
    print(f"  Also:  {{train,val,test}}.parquet, {{train,val,test}}.csv")
    print(f"  Manifest: manifest.json")


def main():
    parser = argparse.ArgumentParser(description="VibeCheck Dataset Preparation Pipeline")
    parser.add_argument("--ember-zip", type=str, default=None,
                        help="Path to EMBER2024 archive.zip (default: datasets/raw/ember2024/archive.zip)")
    parser.add_argument("--max-per-week", type=int, default=5000,
                        help="Max samples to load per weekly JSONL file (default: 5000)")
    parser.add_argument("--val-weeks", type=int, default=8,
                        help="Number of trailing train weeks to use as validation (default: 8)")
    parser.add_argument("--max-benign", type=int, default=50000,
                        help="Max benign PE samples to collect from system directories")
    parser.add_argument("--synthetic-count", type=int, default=0,
                        help="Synthetic vibeware samples to append (default: 0 when using EMBER)")
    parser.add_argument("--bodmas-dir", type=str, default=None,
                        help="Path to BODMAS dataset directory")
    parser.add_argument("--output-dir", type=str, default=None,
                        help="Output directory for processed splits")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--skip-benign-scan", action="store_true",
                        help="Skip scanning system directories for benign samples")
    args = parser.parse_args()

    ensure_dirs()
    output_dir = Path(args.output_dir) if args.output_dir else PROCESSED_DIR
    output_dir.mkdir(parents=True, exist_ok=True)

    zip_path = Path(args.ember_zip) if args.ember_zip else RAW_DIR / "ember2024" / "archive.zip"

    ember_result = load_ember2024_from_zip(
        zip_path,
        split="train",
        max_per_week=args.max_per_week,
        val_weeks=args.val_weeks,
    )

    if ember_result is None:
        print("ERROR: EMBER2024 zip not found or failed to load.")
        print(f"  Expected: {zip_path}")
        return

    X_train, y_train, X_val, y_val = ember_result

    print(f"\n[2] Loading EMBER2024 test split...")
    test_result = load_ember2024_from_zip(
        zip_path,
        split="test",
        max_per_week=args.max_per_week,
        val_weeks=0,
    )
    if test_result is not None:
        X_test, y_test = test_result[0], test_result[1]
    else:
        print("  No test split found, using 15% of train as test.")
        from sklearn.model_selection import train_test_split as tts
        X_train, X_test, y_train, y_test = tts(
            X_train, y_train, test_size=0.15, random_state=args.seed, stratify=y_train
        )

    if not args.skip_benign_scan:
        print(f"\n[3] Collecting benign system PEs (up to {args.max_benign})...")
        benign_paths = collect_benign_samples(max_samples=args.max_benign)
        if benign_paths:
            X_sys, y_sys = extract_pe_features(benign_paths, label=0)
            if X_sys:
                X_sys_arr = np.array(X_sys, dtype=np.float64)
                X_train = np.vstack([X_train, X_sys_arr])
                y_train = np.concatenate([y_train, np.zeros(len(X_sys_arr), dtype=np.int64)])
                print(f"  Added {len(X_sys_arr)} real benign PEs to training set")
    else:
        print("\n[3] Skipping benign system scan")

    if args.bodmas_dir:
        bodmas_dir = Path(args.bodmas_dir)
        if bodmas_dir.exists():
            print(f"\n  Loading BODMAS from {bodmas_dir}...")
            bodmas_files = list(bodmas_dir.rglob("*.exe")) + list(bodmas_dir.rglob("*.dll"))
            if bodmas_files:
                X_b, y_b = extract_pe_features(bodmas_files, label=1)
                if X_b:
                    X_train = np.vstack([X_train, np.array(X_b, dtype=np.float64)])
                    y_train = np.concatenate([y_train, np.array(y_b, dtype=np.int64)])

    if args.synthetic_count > 0:
        print(f"\n[4] Generating {args.synthetic_count} synthetic vibeware samples...")
        X_synth, y_synth = generate_synthetic_vibeware(count=args.synthetic_count, seed=args.seed)
        X_train = np.vstack([X_train, X_synth])
        y_train = np.concatenate([y_train, y_synth])

    print(f"\n  Final splits:")
    for name, X, y in [("train", X_train, y_train), ("val", X_val, y_val), ("test", X_test, y_test)]:
        b = int(np.sum(y == 0))
        m = int(np.sum(y == 1))
        print(f"  {name:5s}: {len(y):7d} samples (benign={b}, malicious={m})")
    spw = np.sum(y_train == 0) / max(np.sum(y_train == 1), 1)
    print(f"  scale_pos_weight = {spw:.4f}")

    for name, X, y in [("train", X_train, y_train), ("val", X_val, y_val), ("test", X_test, y_test)]:
        np.save(output_dir / f"X_{name}.npy", X)
        np.save(output_dir / f"y_{name}.npy", y)
        df = pd.DataFrame(X, columns=FEATURE_NAMES)
        df["label"] = y
        df.to_parquet(output_dir / f"{name}.parquet", index=False)

    manifest = {
        "total_samples": int(len(y_train) + len(y_val) + len(y_test)),
        "train_samples": int(len(y_train)),
        "val_samples": int(len(y_val)),
        "test_samples": int(len(y_test)),
        "feature_names": FEATURE_NAMES,
        "num_features": len(FEATURE_NAMES),
        "source": "EMBER2024",
    }
    with open(output_dir / "manifest.json", "w") as f:
        json.dump(manifest, f, indent=2)

    print(f"\n  Dataset saved to {output_dir}")


if __name__ == "__main__":
    main()