stage2-gnn-pypi-malware

Stage 2 GNN — PyPI Supply-Chain Malware Detection

A GraphSAGE classifier over AST-derived API call graphs, used as the Stage 2 component of a two-stage PyPI supply-chain malware detection pipeline. Stage 1 narrows 800K+ packages to a candidate set via metadata anomaly detection (IsolationForest over graph/metadata features); Stage 2 — this model — scores the candidates' source code.

Model

Input is a directed graph whose nodes are API calls extracted from a package's Python AST and whose edges encode call-sequence and containment relations. Dynamic/unresolvable receivers are normalized to a <dynamic> prefix (e.g. <dynamic>.decode.strip), which keeps the vocabulary bounded across arbitrary third-party code.

Training

Trained on the MalwareBench PyPI corpus, restricted to setup.py and __init__.py — the install-time entry points where supply-chain payloads overwhelmingly live.

• 3,662 malicious / 3,662 benign packages (balanced)
• Max 6 files per package
• 60/20/20 train/val/test split, seed 42
• 40 epochs max, early stopping patience 8, Adam @ lr 1e-3, batch size 16
• Benign set includes mined hard negatives (benign packages with install-time side effects, subprocess use, and network calls)

Evaluation

Held-out test split: 1,573 packages (767 malicious / 806 benign).

• AUROC 0.983
• PR-AUC 0.986

Best F1 is 0.960 at threshold 0.740.

Read these numbers carefully

The test split is balanced 50/50, and real PyPI is not — malicious packages are a very small fraction of uploads. Precision on a balanced split is not precision in deployment. At a realistic base rate the same score distribution yields substantially more false positives per true positive, so the operating threshold has to be re-tuned against the deployed candidate distribution rather than copied from the table above.

The model is also trained and evaluated on a single corpus (MalwareBench), so these figures reflect that corpus's notion of malicious. Generalization to novel campaigns, obfuscation styles, or non-install-time payloads is not measured here.

This checkpoint is one component of a larger pipeline and is not a standalone verdict source. Treat its output as a ranking signal for human review, not as an automated takedown trigger.

Usage

import torch

ckpt = torch.load("stage2_gnn_setup_init_v2.pt", map_location="cpu")

ckpt["config"]      # {'vocab_size': 21085, 'embed_dim': 32, 'hidden': 64, 'dropout': 0.3}
ckpt["vocab"]       # list[str], index -> API-call token ('<unk>' at 0)
ckpt["state_dict"]  # GraphSAGE weights
ckpt["extra"]       # test_report, training_history, training args

Instantiate the matching model definition from the training repository, then model.load_state_dict(ckpt["state_dict"]). Graphs must be built with the same AST extractor and token normalization used in training — the vocabulary is extractor-specific and will not transfer to a differently-built graph.

License

Not yet specified. No reuse rights are granted at this time — please contact the author before redistributing or building on these weights.

Intended use

Security research and defensive triage of the PyPI ecosystem: prioritizing which packages a human analyst looks at. Not intended for automated enforcement, and not evaluated for use outside PyPI Python packages.