H²GLM

H²GLM (Hierarchical Heterogeneous Graph Learning framework enhanced by LLMs for Malicious package detection) is a learning-based framework for detecting malicious Python packages in the software supply chain.

Overview

H²GLM treats a Python package as a hierarchical heterogeneous code graph rather than a flat set of features. By explicitly modeling package-, file-, and function-level entities together with heterogeneous dependencies (containment, import, and function-call relations), and leveraging LLMs to infer function semantics, the framework captures malicious behavior propagation across levels and achieves accurate, robust package-level classification. A function-level attribution mechanism further supports fine-grained localization of malicious behaviors without requiring human expert intervention.

Key Innovations

1. Hierarchical Heterogeneous Graph: Constructs a three-level graph (Package → File → Function) from ASTs, preserving the inherent structural relationships of the project.

2. LLM-Augmented Semantic Node Types (φ^LLM): An LLM dynamically classifies each function into fine-grained behavioral categories (e.g., mal_c2, network, crypto), introducing a second heterogeneity layer that guides the GNN toward semantically meaningful subgraphs.

3. CH²GNN: A custom Heterogeneous Graph Neural Network with φ(v)-guided projection, φ^LLM(v)-guided processing, and ψ(e)-guided message passing, designed specifically for malicious code detection.

4. DropNode Localization: Removes function nodes one at a time and measures the predicted malicious probability drop to identify the functions most strongly associated with malicious behavior.

Architecture

┌──────────────────────────────────────────────────────────────────┐
│                         H²GLM Framework                         │
├──────────────────────────────────────────────────────────────────┤
│  Stage 1: Graph Extraction                                       │
│  ┌─────────────┐    ┌─────────────┐    ┌─────────────┐          │
│  │   Package   │───▶│    File     │───▶│  Function   │          │
│  │    Node     │    │    Node     │    │    Node     │          │
│  └─────────────┘    └─────────────┘    └─────────────┘          │
│        │                  │                   │                  │
│        └──── contains ────┴──── contains ─────┘                  │
│                                               │                  │
│                                    imports / calls               │
├──────────────────────────────────────────────────────────────────┤
│  Stage 2: Hierarchical Heterogeneous Feature Graph Generation    │
│  ┌──────────────────────────────────────────────────────────┐   │
│  │  LLM → φ^LLM(v) function type + semantic description     │   │
│  │  BERT (text) + CodeBERT (code) → node feature vectors    │   │
│  └──────────────────────────────────────────────────────────┘   │
├──────────────────────────────────────────────────────────────────┤
│  Stage 3: Classification (CH²GNN)                                │
│  ┌──────────────────────────────────────────────────────────┐   │
│  │  φ(v)-Guided Projection → φ^LLM(v)-Guided Processing     │   │
│  │  → ψ(e)-Guided Message Passing → READOUT → MLP           │   │
│  └──────────────────────────────────────────────────────────┘   │
├──────────────────────────────────────────────────────────────────┤
│  Stage 4: Locating (DropNode Attribution)                        │
│  ┌──────────────────────────────────────────────────────────┐   │
│  │  Remove v_f → measure ΔP(malicious) → rank & report      │   │
│  └──────────────────────────────────────────────────────────┘   │
└──────────────────────────────────────────────────────────────────┘

Project Structure

H2GLM/
├── README.md
├── requirements.txt
├── src/
│   ├── __init__.py
│   ├── config.py              # Configuration and hyperparameters
│   ├── models.py              # CH²GNN model definition
│   ├── preprocess.py          # Graph preprocessing and feature extraction
│   ├── data_extraction/
│   │   ├── __init__.py
│   │   ├── batch_analyze.py   # AST-based package analysis
│   │   └── enrich_json.py     # LLM-based semantic enrichment (φ^LLM)
│   └── predict.py             # Inference pipeline

Dataset

The malicious PyPI package dataset (2,134 samples) used in our evaluation is publicly available on Zenodo:

DOI: 10.5281/zenodo.20814365

Download and extract to data/:

wget https://zenodo.org/records/20814365/files/malicious_2134.tar.gz
tar -xzf malicious_2134.tar.gz -C data/

Installation

git clone https://github.com/xxy33/malware.git
cd malware

pip install -r requirements.txt

Requirements

• Python >= 3.8
• PyTorch >= 1.12
• PyTorch Geometric
• Transformers
• Sentence-Transformers
• OpenAI (for LLM enrichment)

Usage

1. Graph Extraction

Extract the heterogeneous code graph from Python packages:

python -m src.data_extraction.batch_analyze \
    -i /path/to/packages \
    -o /path/to/output \
    -w /path/to/whitelist.txt

2. LLM Semantic Enrichment (φ^LLM)

Classify each function into behavioral node types and generate semantic descriptions:

python -m src.data_extraction.enrich_json \
    -d /path/to/analysis_results

The LLM prompt templates and the full list of function-type categories are documented in this README (see Function Content Types below).

3. Feature Extraction

Convert JSON graphs to PyTorch Geometric format:

python -m src.preprocess

4. Prediction

Run inference on new packages:

python -m src.predict \
    -i /path/to/package_list.txt \
    -o /path/to/results.json

Model Architecture

Node Types

Type	Description
Package	Root node representing the entire software package
File	Python source files within the package
Function	Functions and methods extracted from source files

Edge Types

Type	Direction	Description
contains	Package→File, File→Function	Hierarchical containment
imports	File→File	Inter-file import dependencies
calls	Function→Function	Function call relationships

Feature Dimensions

Component	Dimension
Text embeddings (Sentence Transformer)	768
Code embeddings (CodeBERT)	768
Aligned feature dimension	256
CH²GNN hidden dimension	256
CH²GNN output dimension	128

Function Content Types

H²GLM uses an LLM to dynamically classify each function into one of the following behavioral categories (φ^LLM). These are derived by running a centralized LLM scan across the target dataset rather than being pre-defined labels.

Type	Description
setup	Installation scripts and initialization
network	HTTP requests, socket operations, outbound communication
file	Local file I/O operations
process	External command execution
crypto	Encryption, hashing, and key operations
data	Data processing and transformation
util	Helper utilities
mal_downloader	Malicious download-and-execute patterns
mal_c2	Command & Control communication or credential exfiltration
unknown	Unclassified functions

Configuration

Key hyperparameters in src/config.py:

Parameter	Default	Description
BATCH_SIZE	16	Training batch size
LEARNING_RATE	0.001	Learning rate
EPOCHS	100	Number of training epochs
HGT_HEADS	4	Number of attention heads
DROPOUT_RATE	0.3	Dropout rate

License

This project is licensed under the MIT License — see the LICENSE file for details.

Acknowledgments

• PyTorch Geometric for heterogeneous graph neural network primitives
• Hugging Face Transformers for CodeBERT and Sentence Transformer models
• OpenAI / Azure OpenAI for LLM-based semantic enrichment