predict.py

"""
predict.py — ViGil Standalone Malware Predictor

Run inference with the trained JointMalwareModel on a single file.

Architecture matches traning_notebook/vigil.ipynb exactly:
  OptimizedHGT   (CPG graph,  HIDDEN=384, LAYERS=6) → 768-dim
  OptimizedCNN   (ConvNeXt-Tiny grayscale image)     → 512-dim
  OptimizedRansomFormerEncoder (bytes + API imports) → 256-dim
  OptimizedFusion (Deep Residual MLP)    fused 1536  → 2-class

Usage:
    python predict.py --file suspicious.exe
    python predict.py --file sample.dll --model joint_model.pt --samples 30
    python predict.py --file document.pdf --verbose
    python predict.py --file sample.exe --json

Requirements:
    pip install torch torchvision pillow numpy
"""

import argparse
import sys
import json
import logging
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("vigil.predict")

# ── Default paths ─────────────────────────────────────────────────────────────
_PROJECT_ROOT  = Path(__file__).resolve().parent
_DEFAULT_MODEL = _PROJECT_ROOT / "joint_model.pt"
_MODEL_CONFIG  = _PROJECT_ROOT / "model_config.json"

# ── Canonical architecture config (notebook values, used as safe fallback) ────
_NOTEBOOK_CFG = {
    "embedding_dim":  320,
    "hidden_dim":     384,
    "num_heads":      8,
    "num_layers":     6,
    "num_classes":    2,
    "fused_dim":      1536,   # HGT(768) + CNN(512) + RF(256)
    "byte_seq_len":   1024,
    "max_apis":       256,
    "api_vocab_size": 4096,
    "label_map":      {"0": "BENIGN", "1": "MALWARE"},
}


def _load_model_cfg(checkpoint_path: Path = None) -> dict:
    """Read model_config.json if available, fall back to notebook defaults."""
    cfg = {**_NOTEBOOK_CFG}
    if _MODEL_CONFIG.exists():
        with open(_MODEL_CONFIG) as fh:
            on_disk = json.load(fh)
        cfg.update(on_disk)

    if checkpoint_path and checkpoint_path.exists():
        try:
            import torch
            ckpt = torch.load(checkpoint_path, map_location="cpu")
            state = ckpt.get("model_state", ckpt)
            if "hgt.proj.0.weight" in state:
                weight_shape = state["hgt.proj.0.weight"].shape
                if len(weight_shape) == 2:
                    detected_dim = weight_shape[1]
                    if detected_dim != cfg.get("embedding_dim"):
                        logger.info(f"Auto-detected embedding_dim={detected_dim} from checkpoint (was {cfg.get('embedding_dim')})")
                        cfg["embedding_dim"] = detected_dim
        except Exception as e:
            logger.warning(f"Could not auto-detect embedding_dim from checkpoint: {e}")
    return cfg


def _build_model(cfg: dict, device: "torch.device"):
    """Reconstruct the exact notebook architecture."""
    from uir.model.optimized_models import build_model
    return build_model(cfg, device)


def _load_checkpoint(model, checkpoint_path: Path, device):
    """Load model weights from a .pt checkpoint."""
    import torch

    if not checkpoint_path.exists():
        raise FileNotFoundError(
            f"Checkpoint not found: {checkpoint_path}\n"
            "Train the model first via traning_notebook/vigil.ipynb on Kaggle, "
            "then place the checkpoint at models/01/models/joint_model.pt"
        )

    logger.info(f"Loading checkpoint: {checkpoint_path}")
    ckpt  = torch.load(checkpoint_path, map_location=device)
    state = ckpt.get("model_state", ckpt)   # handle both wrapped and bare state_dict
    model.load_state_dict(state, strict=False)
    model.eval()
    logger.info("Checkpoint loaded.")
    return model


def _extract_features(file_path: Path, cfg: dict, device, verbose: bool):
    """
    Extract all four modality features from a single file.

    Returns a namespace with:
        .x, .edge_index, .node_types, .edge_types  — CPG / HGT inputs
        .image                                      — [3, 224, 224]
        .pe_bytes                                   — [1, 1024]
        .api_tokens                                 — [max_apis]
    """
    import torch
    import torchvision.transforms as T
    from uir.pipeline.processor import FileProcessor
    from uir.config import UIRConfig
    from uir.model.dataset import CPGDataset
    from uir.extraction.image_generator import pe_to_grayscale_image
    from uir.extraction.pe_feature_extractor import extract_ransomformer_features

    config    = UIRConfig()
    processor = FileProcessor(config)

    # ── 1. CPG extraction ─────────────────────────────────────────────────────
    if verbose:
        logger.info("Extracting Code Property Graph …")
    cpg = processor.process(file_path, use_cache=False)
    if cpg is None:
        raise RuntimeError(
            f"CPG extraction failed for {file_path}. "
            "Ensure the uir package and its dependencies are installed."
        )

    # Convert CPG → tensor data (node features, edge index, etc.)
    _dummy = CPGDataset(cpg_dir=file_path.parent, embedding_dim=cfg.get("embedding_dim", 320))
    data   = _dummy._cpg_to_data(cpg)

    # ── 2. Grayscale byte-image ───────────────────────────────────────────────
    if verbose:
        logger.info("Generating grayscale byte-image …")
    img    = pe_to_grayscale_image(file_path, target_size=224).convert("RGB")
    img_tf = T.Compose([
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225]),
    ])
    data.image = img_tf(img)                              # [3, 224, 224]

    # ── 3. RansomFormer inputs ────────────────────────────────────────────────
    if verbose:
        logger.info("Extracting byte sequence and API import tokens …")
    api_names = list(cpg.metadata.get("imports", [])) if cpg.metadata else []
    pe_bytes, api_tokens = extract_ransomformer_features(
        file_path,
        api_names = api_names,
        seq_len   = cfg.get("byte_seq_len", 1024),
        max_apis  = cfg.get("max_apis",      256),
    )
    data.pe_bytes   = pe_bytes    # [1, 1024]
    data.api_tokens = api_tokens  # [max_apis]

    return data


def _collate_single(data, device):
    """Wrap a single CPGData into a batch-1 tuple for the model."""
    from uir.model.dataset import collate_cpg_batch
    batch_data, batch_idx = collate_cpg_batch([data])
    batch_data = batch_data.to(device)
    batch_idx  = batch_idx.to(device)
    return batch_data, batch_idx


def predict(
    file_path: Path,
    checkpoint_path: Path = _DEFAULT_MODEL,
    num_samples: int = 20,
    verbose: bool = False,
    device_str: str = "auto",
) -> dict:
    """
    Run quad-modal BNN inference on a single file.

    Args:
        file_path:       Path to the file to analyse.
        checkpoint_path: Path to joint_model.pt checkpoint.
        num_samples:     Monte Carlo dropout sampling iterations.
        verbose:         Log intermediate steps.
        device_str:      'auto' | 'cpu' | 'cuda' | 'mps'.

    Returns:
        dict with keys: file, prediction, label, confidence, variance
    """
    import torch

    # ── Device selection ──────────────────────────────────────────────────────
    if device_str == "auto":
        if torch.cuda.is_available():
            device = torch.device("cuda")
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            device = torch.device("mps")
        else:
            device = torch.device("cpu")
    else:
        device = torch.device(device_str)

    logger.info(f"Device: {device}")

    # ── Config + model ────────────────────────────────────────────────────────
    cfg   = _load_model_cfg(checkpoint_path)
    model = _build_model(cfg, device)
    model = _load_checkpoint(model, checkpoint_path, device)

    # ── Validate input file ───────────────────────────────────────────────────
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"Target file not found: {file_path}")

    # ── Feature extraction ────────────────────────────────────────────────────
    data = _extract_features(file_path, cfg, device, verbose)
    batch_data, batch_idx = _collate_single(data, device)

    # ── Monte Carlo inference ─────────────────────────────────────────────────
    if verbose:
        logger.info(f"Running Monte Carlo inference (T={num_samples} samples) …")

    preds, confidence, variance = model.predict_with_confidence(
        batch_data.x,          batch_data.edge_index,
        batch_data.node_types, batch_data.edge_types,
        batch_idx,             batch_data.image,
        batch_data.pe_bytes,   batch_data.api_tokens,
        num_samples=num_samples,
    )

    label_map  = cfg.get("label_map", {"0": "BENIGN", "1": "MALWARE"})
    pred_class = preds[0].item()
    label      = label_map.get(str(pred_class), "UNKNOWN")
    conf       = confidence[0].item()
    var        = variance[0].item()

    return {
        "file":       str(file_path),
        "prediction": pred_class,
        "label":      label,
        "confidence": conf,
        "variance":   var,
    }


def _print_result(result: dict):
    """Pretty-print the prediction result to stdout."""
    label = result["label"]
    conf  = result["confidence"] * 100
    var   = result["variance"]

    try:
        import os
        supports_color = os.isatty(sys.stdout.fileno())
    except Exception:
        supports_color = False

    G = "\033[92m"; R = "\033[91m"; Y = "\033[93m"; RESET = "\033[0m"
    if not supports_color:
        G = R = Y = RESET = ""

    label_col = f"{G}{label}{RESET}" if label == "BENIGN" else f"{R}{label}{RESET}"

    border = "=" * 66
    print(f"\n{border}")
    print("  ViGil — Quad-Modal Malware Detection")
    print("  Architecture: OptimizedHGT + ConvNeXt + AttentionByte + CLSTransformerAPI → DeepResMLP")
    print(border)
    print(f"  File:        {Path(result['file']).name}")
    print(f"  Verdict:     {label_col}")
    print(f"  Confidence:  {Y}{conf:.2f}%{RESET}")
    print(f"  Uncertainty: {var:.6f}  (epistemic variance, MC dropout)")
    print(f"{border}\n")


def main():
    parser = argparse.ArgumentParser(
        description=(
            "ViGil — Quad-Modal Malware Detector\n"
            "Architecture: OptimizedHGT + ConvNeXt-Tiny + AttentionByteEncoder"
            " + CLSTransformerAPI → DeepResMLP"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python predict.py --file suspicious.exe
  python predict.py --file sample.dll --samples 30 --verbose
  python predict.py --file doc.pdf --model models/01/models/joint_model.pt --json
        """,
    )
    parser.add_argument("--file",    "-f", required=True,
                        help="Path to the file to analyse.")
    parser.add_argument("--model",   "-m", default=str(_DEFAULT_MODEL),
                        help=f"Checkpoint path (default: {_DEFAULT_MODEL}).")
    parser.add_argument("--samples", "-s", type=int, default=20,
                        help="Monte Carlo dropout samples (default: 20).")
    parser.add_argument("--device",  default="auto",
                        choices=["auto", "cpu", "cuda", "mps"],
                        help="Compute device (default: auto).")
    parser.add_argument("--verbose", "-v", action="store_true",
                        help="Verbose progress logging.")
    parser.add_argument("--json",    action="store_true",
                        help="Output result as JSON.")

    args = parser.parse_args()

    if not args.verbose:
        logging.getLogger().setLevel(logging.WARNING)

    try:
        result = predict(
            file_path       = Path(args.file),
            checkpoint_path = Path(args.model),
            num_samples     = args.samples,
            verbose         = args.verbose,
            device_str      = args.device,
        )
    except FileNotFoundError as exc:
        print(f"\n[ERROR] {exc}", file=sys.stderr)
        sys.exit(1)
    except Exception as exc:
        print(f"\n[ERROR] Prediction failed: {exc}", file=sys.stderr)
        if args.verbose:
            import traceback
            traceback.print_exc()
        sys.exit(1)

    if args.json:
        print(json.dumps(result, indent=2))
    else:
        _print_result(result)


if __name__ == "__main__":
    main()