You can deploy this model and start earning money today!

Azaz666/
stage2-gnn-pypi-malware

Name: stage2-gnn-pypi-malware
Availability: InStock

Views

No views yet

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


Architecture	2-layer GraphSAGE + 2-layer MLP classifier
Embedding	32-dim, initialized from Word2Vec over API-call sequences
Hidden dim	64
Dropout	0.3
Vocabulary	21,085 API-call tokens
Parameters	~840K
Checkpoint size	3.4 MB

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

Threshold	Precision	Recall	F1	Accuracy
0.30	0.975	0.932	0.953	0.955
0.50	0.985	0.931	0.957	0.959
0.70	0.993	0.930	0.960	0.962
0.85	0.993	0.922	0.956	0.959

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

python

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.

Azaz666/stage2-gnn-pypi-malware