train_malware_model.py

#!/usr/bin/env python3
"""Train malware one-class GNN model using malware samples only."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import json
import os
from pathlib import Path

import torch

from dataset import MalwareGraphDataset, merge_manifest_csv_files
from analysis_schema import SCHEMA_VERSION
from one_class_gnn import train_one_class

DEFAULT_ALLOWED_BENIGN_SUBTYPES = (
    "clean_benign,hard_benign_admin_tooling,ambiguous_novirus_control"
)


def resolve_allowed_benign_subtypes(args) -> tuple[str, ...] | None:
    if not args.require_governance_manifest:
        return None
    return tuple(
        x.strip()
        for x in str(args.allowed_benign_subtypes).split(",")
        if x.strip()
    )


def run(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[MalwareModel] Device: {device}")
    allowed_benign_subtypes = resolve_allowed_benign_subtypes(args)
    print(f"[MalwareModel] Governance manifest required: {bool(args.require_governance_manifest)}")
    if allowed_benign_subtypes is not None:
        print(
            "[MalwareModel] Benign subtype allowlist (label=0 rows):",
            ", ".join(allowed_benign_subtypes),
        )
    print(
        "[MalwareModel] Strict train filter:",
        not args.disable_strict_train_filter,
    )

    manifest_path = args.manifest
    merged_tmp = None
    extra_manifests = [x.strip() for x in str(getattr(args, "extra_manifest", "") or "").split(",") if x.strip()]
    if extra_manifests:
        manifest_path = merge_manifest_csv_files(args.manifest, extra_manifests)
        merged_tmp = manifest_path
        print(f"[MalwareModel] Merged {len(extra_manifests)} extra manifest(s) → {manifest_path}")
    try:
        ds = MalwareGraphDataset(
            manifest_path,
            base_dir=args.base_dir,
            include_uncertain=not args.exclude_uncertain,
            include_unknown=False,
            require_train_eligible=not args.disable_strict_train_filter,
            require_governance_columns=args.require_governance_manifest,
            allowed_benign_subtypes=allowed_benign_subtypes,
            target="label",
            graph_attr_profile=getattr(args, "graph_attr_profile", "full"),
            graph_view=getattr(args, "graph_view", "full"),
        )
    finally:
        if merged_tmp and merged_tmp != os.path.abspath(args.manifest) and os.path.isfile(merged_tmp):
            try:
                os.remove(merged_tmp)
            except OSError:
                pass
    malware_ds = [ds[i] for i in range(len(ds)) if int(ds[i].y.item()) == 1]
    if len(malware_ds) < 10:
        raise SystemExit("[ERROR] Need at least 10 malware samples for one-class malware model")
    print(f"[MalwareModel] Malware training set: {len(malware_ds)} graph(s)")

    art = train_one_class(
        malware_ds,
        device,
        hidden=args.hidden,
        layers=args.layers,
        dropout=args.dropout,
        edge_emb_dim=args.edge_emb_dim,
        out_dim=args.out_dim,
        epochs=args.epochs,
        batch_size=args.batch_size,
        lr=args.lr,
        weight_decay=args.weight_decay,
        radius_quantile=args.radius_quantile,
        seed=args.seed,
        contrastive_weight=args.contrastive_weight,
        radius_min=args.radius_min,
        radius_max=args.radius_max,
        trim_fraction=args.trim_fraction,
        calibration_fraction=args.calibration_fraction,
        center_update_interval=args.center_update_interval,
    )

    payload = {
        "model_type": "malware_pattern_one_class_gnn",
        "state_dict": art.state_dict,
        "center": art.center,
        "radius": art.radius,
        "in_channels": art.in_channels,
        "graph_attr_dim": art.graph_attr_dim,
        "hidden": art.hidden,
        "layers": art.layers,
        "dropout": art.dropout,
        "edge_emb_dim": art.edge_emb_dim,
        "out_dim": art.out_dim,
        "train_size": art.train_size,
        "calibration_size": art.calibration_size,
        "radius_source": art.radius_source,
        "radius_quantile": art.radius_quantile,
        "trim_fraction": art.trim_fraction,
        "calibration_fraction": art.calibration_fraction,
        "center_update_interval": art.center_update_interval,
        "score_threshold_low": art.score_threshold_low,
        "score_threshold_high": art.score_threshold_high,
        "schema_version": SCHEMA_VERSION,
        "trained_at_utc": datetime.now(timezone.utc).isoformat(),
    }
    out_model = Path(args.output_model)
    out_model.parent.mkdir(parents=True, exist_ok=True)
    torch.save(payload, out_model)

    meta = {
        "model_type": payload["model_type"],
        "manifest": args.manifest,
        "extra_manifest_csvs": extra_manifests,
        "base_dir": args.base_dir,
        "exclude_uncertain": args.exclude_uncertain,
        "strict_train_filter": (not args.disable_strict_train_filter),
        "require_governance_manifest": bool(args.require_governance_manifest),
        "allowed_benign_subtypes": list(allowed_benign_subtypes or ()),
        "n_malware_train": len(malware_ds),
        "seed": args.seed,
        "hidden": args.hidden,
        "layers": args.layers,
        "dropout": args.dropout,
        "edge_emb_dim": args.edge_emb_dim,
        "out_dim": args.out_dim,
        "epochs": args.epochs,
        "batch_size": args.batch_size,
        "lr": args.lr,
        "weight_decay": args.weight_decay,
        "radius_quantile": args.radius_quantile,
        "radius_min": args.radius_min,
        "radius_max": args.radius_max,
        "contrastive_weight": args.contrastive_weight,
        "trim_fraction": args.trim_fraction,
        "calibration_fraction": args.calibration_fraction,
        "center_update_interval": args.center_update_interval,
        "radius": art.radius,
        "radius_source": art.radius_source,
        "train_size": art.train_size,
        "calibration_size": art.calibration_size,
        "score_threshold_low": art.score_threshold_low,
        "score_threshold_high": art.score_threshold_high,
        "schema_version": SCHEMA_VERSION,
        "trained_at_utc": datetime.now(timezone.utc).isoformat(),
    }
    out_meta = Path(args.output_meta)
    out_meta.parent.mkdir(parents=True, exist_ok=True)
    out_meta.write_text(json.dumps(meta, indent=2), encoding="utf-8")
    print(f"[MalwareModel] saved: {out_model}")
    print(f"[MalwareModel] meta : {out_meta}")


if __name__ == "__main__":
    p = argparse.ArgumentParser(description="Train malware one-class GNN model")
    p.add_argument("manifest", help="Path to dataset_manifest.csv")
    p.add_argument("--base-dir", default=None, dest="base_dir")
    p.add_argument("--exclude-uncertain", action="store_true", dest="exclude_uncertain")
    p.add_argument(
        "--disable-strict-train-filter",
        action="store_true",
        dest="disable_strict_train_filter",
        help="Allow rows with train_eligible=false (default strict filter keeps them out).",
    )
    p.add_argument(
        "--require-governance-manifest/--no-require-governance-manifest",
        default=True,
        action=argparse.BooleanOptionalAction,
        dest="require_governance_manifest",
        help="Require train_eligible and related governance columns in the manifest.",
    )
    p.add_argument(
        "--allowed-benign-subtypes",
        default=DEFAULT_ALLOWED_BENIGN_SUBTYPES,
        dest="allowed_benign_subtypes",
        help=(
            "Comma-separated benign_subtype values permitted when governance is on. "
            "Applies to label=0 rows during manifest load; malware (label=1) rows are "
            "selected separately. train_eligible=true bypasses this list when strict "
            "train filter is on."
        ),
    )
    p.add_argument("--hidden", type=int, default=32)
    p.add_argument("--layers", type=int, default=2)
    p.add_argument("--dropout", type=float, default=0.4)
    p.add_argument("--edge-emb-dim", type=int, default=16, dest="edge_emb_dim")
    p.add_argument("--out-dim", type=int, default=64, dest="out_dim")
    p.add_argument("--epochs", type=int, default=120)
    p.add_argument("--batch-size", type=int, default=4, dest="batch_size")
    p.add_argument("--lr", type=float, default=3e-4)
    p.add_argument("--weight-decay", type=float, default=1e-4, dest="weight_decay")
    p.add_argument("--radius-quantile", type=float, default=0.9, dest="radius_quantile")
    p.add_argument("--radius-min", type=float, default=1e-3, dest="radius_min")
    p.add_argument("--radius-max", type=float, default=25.0, dest="radius_max")
    p.add_argument("--trim-fraction", type=float, default=0.10, dest="trim_fraction")
    p.add_argument("--calibration-fraction", type=float, default=0.20, dest="calibration_fraction")
    p.add_argument("--center-update-interval", type=int, default=5, dest="center_update_interval")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--contrastive-weight", type=float, default=0.10, dest="contrastive_weight")
    p.add_argument("--output-model", default="outputs/malware_model.pt", dest="output_model")
    p.add_argument("--output-meta", default="outputs/malware_model_meta.json", dest="output_meta")
    p.add_argument(
        "--extra-manifest",
        default="",
        dest="extra_manifest",
        help="Comma-separated extra manifest CSV paths merged after primary (dedupe by folder; primary wins).",
    )
    p.add_argument(
        "--graph-attr-profile",
        choices=["full", "no_manifest_leakage", "structure_only"],
        default="no_manifest_leakage",
        dest="graph_attr_profile",
    )
    p.add_argument(
        "--graph-view",
        choices=["full", "attack_subgraph"],
        default="full",
        dest="graph_view",
    )
    run(p.parse_args())