LLM-Enhanced Hierarchical Heterogeneous Graph Representation Learning for Malicious Python Package Detection

Technical Summary: H2GLM for Malicious Python Package Detection

Problem Statement
The widespread adoption of open-source repositories like PyPI has made malicious Python packages a critical threat to software supply chains. Existing detection methods face significant limitations:

• Rule-based approaches (e.g., signatures, heuristics) struggle with obfuscation, dynamic code loading, and rapidly evolving attack patterns, leading to high false-positive and false-negative rates.
• Learning-based approaches often fail to capture the hierarchical organization and heterogeneous interactions inherent in Python packages. While Graph Neural Networks (GNNs) model structural dependencies, they frequently rely on homogeneous graphs that treat all entities uniformly, missing critical distinctions between packages, files, and functions.
• Large Language Models (LLMs) excel at semantic reasoning but lack explicit structural context and struggle to process large-scale, interconnected program structures without truncation, leading to incomplete attack-chain reconstruction.
• Crucially, malicious behaviors often emerge from interactions across heterogeneous components (e.g., installation scripts invoking functions that access environment variables and exfiltrate data), a complexity that current homogeneous GNN-LLM frameworks cannot effectively model.

Methodology: H2GLM Framework
The authors propose H2GLM (Hierarchical Heterogeneous Graph representation learning framework enhanced by LLMs), a four-stage framework designed to jointly model explicit structural heterogeneity and implicit semantic heterogeneity.

1. Graph Extraction:

   • Constructs a heterogeneous code graph $G = (V, E, \phi, \psi)$ from the Abstract Syntax Tree (AST).
   • Nodes ($V$): Three types: Package, File, and Function.
   • Edges ($E$): Three types: Contain (package-file, file-function), Import, and Call.
   • Attributes: Source code for functions; LLM-generated summaries for files and packages based on project manifests and docstrings.

2. Hierarchical Heterogeneous Feature Graph Generation:

   • Enhances the graph with a novel layer of semantic heterogeneity, $\phi_{LLM}$.
   • LLM Inference: An LLM analyzes the source code of each Function node to dynamically infer its functional semantic type (e.g., environment inspection, credential access, network communication) and generate a semantic description. This introduces latent heterogeneity not present in the raw structure.
   • Feature Encoding: Node features are constructed by concatenating embeddings from a Sentence Transformer (for text summaries) and CodeBERT (for source code snippets), creating rich, type-aware feature vectors.

3. Classification (H2GNN):

   • A novel Hierarchical Heterogeneous Graph Neural Network (H2GNN) processes the graph through three specific mechanisms:
     • $\phi(v)$-Guided Feature Projection: Projects heterogeneous node types (Package, File, Function) into a unified latent space using type-specific learnable matrices.
     • $\phi_{LLM}(v)$-Guided Feature Processing: Further processes node representations based on the LLM-inferred functional types, allowing the model to learn specialized representations for different semantic roles.
     • $\psi(e)$-Guided Message Passing: Aggregates information across edges, using type-specific weights and attention coefficients to emphasize critical relations (e.g., import chains, function calls) indicative of malicious propagation.
   • A global mean pooling operation followed by an MLP predicts the package-level maliciousness.

4. Localization:

   • To identify specific malicious functions, the framework employs a function-level attribution mechanism.
   • It iteratively removes function nodes from the graph and measures the drop in the predicted malicious probability.
   • Functions causing significant probability drops are flagged as suspicious. Their local subgraphs are then analyzed by an LLM to generate human-interpretable explanations of the attack logic.

Key Contributions

• H2GLM Framework: A learning-based framework that integrates GNNs and LLMs to model both explicit structural heterogeneity (packages, files, functions) and implicit semantic heterogeneity (LLM-inferred functional roles).
• Novel H2GNN Architecture: A graph neural network tailored for code analysis that performs type-aware message passing and content-aware feature processing, enabling the modeling of malicious behavior propagation across hierarchical levels.
• Automated Localization: A mechanism that automatically identifies suspicious functions and attack paths without human intervention, combining graph attribution with LLM reasoning for interpretability.
• Real-World Validation: The framework was deployed to monitor PyPI, successfully identifying and reporting malicious packages that were subsequently confirmed and removed by maintainers.

Experimental Results
Extensive experiments were conducted on a dataset of 3,049 malicious and numerous benign packages, stratified by size and import complexity.

• Overall Performance: H2GLM achieved an F1-score of 97.12%, consistently outperforming state-of-the-art baselines including traditional ML methods, graph-based detectors (MPHunter), static analysis tools (Bandit4Mal, OSSGadget), and various LLMs (Qwen3-14b, DeepSeek-R1-70b).
• Robustness: H2GLM maintained stable performance (F1 > 94.9%) across extreme conditions, including extremely small packages (<1.57KB) and packages with extensive import dependencies (>100 imports), where other methods (especially LLMs and homogeneous GNNs) suffered significant performance degradation due to context limits or inability to model complex dependencies.
• Ablation Studies: Removing the heterogeneous message passing ($\psi(e)$) or the LLM-guided semantic processing ($\phi_{LLM}$) resulted in F1-score drops of ~4%, confirming the necessity of both structural and semantic heterogeneity.
• Localization Accuracy: On a manually annotated subset of 300 packages, H2GLM achieved 95.67% accuracy in identifying the specific functions responsible for malicious behavior.
• Case Studies: The system successfully identified obfuscated command-and-control (C2) communication in a malicious package masquerading as a cryptographic library, pinpointing the specific network function despite heavy obfuscation.

Significance and Claims
The paper claims that H2GLM represents a significant advancement in software supply chain security by addressing the "heterogeneity gap" in existing detection methods. By explicitly modeling the multi-level structure of Python packages and leveraging LLMs to uncover latent semantic roles, the framework achieves:

1. Higher Accuracy: Superior detection rates compared to both rule-based and existing learning-based methods.
2. Robustness: Effective performance across diverse package complexities where other methods fail.
3. Interpretability: The ability to not only classify packages but also localize and explain the specific malicious behaviors, reducing the need for manual expert intervention.
4. Practical Applicability: Demonstrated success in a real-world, automated monitoring pipeline on PyPI, validating its utility for continuous large-scale ecosystem protection.

The authors emphasize that this work is the first to integrate hierarchical heterogeneous graph learning with LLMs specifically for malicious package detection, offering a paradigm shift from homogeneous representations to a more nuanced, behavior-driven analysis.