AryanPaneru12/HCPG_GNN_Vulnerability_Detection
GitHub: AryanPaneru12/HCPG_GNN_Vulnerability_Detection
基于 HCPG 图表示与图神经网络的智能合约漏洞检测框架，将 Solidity 代码转化为统一图结构并用 GATv2 分类器识别多函数交互型安全漏洞。
Stars: 0 | Forks: 0
# HCPG-GNN：智能合约漏洞检测框架
一个智能合约漏洞检测框架，它构建 HCPG 风格的图视图，并在有兼容权重时使用 Feature-GATv2 分类器。如果 GNN 运行时或当前的 checkpoint 不可用，API 将回退到确定性的特征和 regex 检查。
**项目：** BCSE498J - 毕业设计
**作者：** Aryan Paneru (22BCE3796), Ayush Chaudhary (22BCE3913)
**机构：** VIT 大学，计算机科学与工程学院
## 架构
```
Solidity source
-> input validation and parser
-> function call graph + control-flow graph + HCPG-style unified graph
-> shared structural feature extractor
-> Feature-GATv2 checkpoint when torch/PyG and compatible weights are present
-> deterministic structural detectors for high-signal patterns
-> API response with findings, graph JSON, pipeline stages, and metrics scope
```
## 功能
- **HCPG 风格的图构建**：结合函数、控制流、调用和数据流边的统一可视化
- **Feature-GATv2 运行时**：当存在兼容的 checkpoint 时，对每个函数的结构特征使用 3 层 GATv2 分类器
- **确定性回退**：当 PyTorch/PyG 或兼容权重不可用时，依然提供特征和 regex 检查
- **归因权重**：兼容的 GNN 运行使用 GAT 注意力权重，回退运行使用基于特征的热力图
- **REST API**：用于程序化访问的 FastAPI 后端
- **Web 界面**：用于分析的交互式仪表板
## 性能
结果从 `models/model_metrics.json` 加载；没有硬编码第三方基准分数。当前提交的 artifact 报告如下：
| 模型 | 精确匹配准确率 | 位准确率 | Precision | Recall | F1 | Macro-F1 | AUC-ROC |
|-------|----------------|--------------|-----------|--------|----|----------|---------|
| Feature-GATv2 (本框架) | 83.87% | 95.16% | 80.00% | 94.62% | 0.867 | 0.580 | 0.907 |
### 为什么这些数字与早期草稿不同
该代码库的早期版本使用合成生成的图（`synthetic_hcpg_dataset`），由于学习到的是可学习的噪声模式而非真正的漏洞信号，从而产生了虚高的分数（约 94-95%）。当前的训练脚本克隆了真实的 SmartBugs Curated 和 SolidiFI-benchmark 仓库，并直接从 `.sol` 源文件中提取结构特征。由此产生的指标较低，但对任何独立审查者来说都是真实且可复现的。
在当前的数据划分中，少数类的支持度仍然很小：Access Control 有 3 个测试样本，Integer Overflow 有 2 个，Delegatecall 有 1 个。在添加更大规模的人工标注数据集之前，请将这些类别的类级指标视为参考方向。
## 数据集
| 数据集 | 提供的内容 | 来源 | 标签质量 |
|---------|------------------|--------|---------------|
| SmartBugs Curated | 少量带有人工标记的公开漏洞 Solidity 合约 | [GitHub](https://github.com/smartbugs/smartbugs-curated) | 人工整理的 SWC/类别标签 |
| SolidiFI-benchmark | 注入了智能合约漏洞的较大基准数据集 | [GitHub](https://github.com/DependableSystemsLab/SolidiFI-benchmark) | 对覆盖率有帮助，但注入法并不能完全匹配真实的审计分布 |
训练脚本（`models/colab_hgt_visual.py`）在运行时克隆这两个仓库，并根据文件夹名称派生出漏洞标签。漏洞文件夹变为六个多标签输出，而 safe/clean 文件夹或 `SAFE_DATASET_DIRS` 中的示例变为全零目标。该 pipeline 会对标准化后的 Solidity 源码进行去重，报告任何标签冲突，使用感知组的训练集/验证集/测试集划分以减少项目系列泄漏，并通过对验证集 Macro-F1 进行早停来选择保存的 checkpoint，从而确保少数类不会被多数类的性能所掩盖。各个类别的指标和阈值会被导出，以防止虚高的位级准确率被误认为是精确的多标签性能。
要添加更强大的人工标注基准数据集，请将其解压到按漏洞标记的文件夹中，并使用 `EXTRA_DATASET_DIRS=/path/to/dataset` 运行训练。要添加经过审查的干净合约，请将它们放在 safe/clean 文件夹中，或传入 `SAFE_DATASET_DIRS=/path/to/safe_contracts`。未知的文件夹标签会被跳过，而不会被强行归入错误的类别。保存的 `model_metrics.json` 记录了 `dataset_sources`、`class_distribution`、`safe_contracts`、经过验证集调优的 `thresholds`、`dedupe_report`、`split_report`、`training_config` 和 `training_history`；`model_dataset_manifest.json` 记录了来源和划分大小；`model_split_manifest.json` 记录了确切的训练集/验证集/测试集成员身份；`best_hgt_model.pt` 包含 checkpoint 元数据；`model_error_analysis.json` 按漏洞类别列出了假阳性和假阴性样本供审查；`calibration_report.json` 记录了可靠性分箱、预期校准误差和 Brier score，用于验证/测试集概率；`calibration_curves.png` 可视化了未校准与温度缩放后的可靠性。
## 项目结构
```
final year project/
+-- frontend/
| +-- index.html # Web interface
+-- backend/
| +-- app.py # FastAPI application
+-- models/
| +-- artifact_validator.py # Artifact compatibility and reproducibility audit
| +-- colab_hgt_visual.py # Main training script (real datasets, shared feature schema)
| +-- experiment_tracking.py # JSONL training run logging
| +-- metrics_utils.py # Threshold calibration helpers
| +-- training_utils.py # Dedupe, grouping, split, and JSON export helpers
| +-- model_metrics.json # Metrics/checkpoint metadata
+-- docs/
| +-- ML_ENGINEERING_AUDIT.md # Implemented/pending ML engineering report
| +-- MODEL_CARD.md # Model scope, metrics, limitations, safety notes
| +-- DATA_CARD.md # Dataset provenance, labels, bias, split strategy
+-- scripts/
| +-- validate_artifacts.py # Local/CI artifact validation CLI
| +-- generate_model_quality_report.py # Human-readable retrain report
| +-- run_ml_release_gate.py # Final ML release readiness gate
| +-- run_model_canary_suite.py # Behavioral API/model canary checks
| +-- run_post_promotion_verification.py # One-command post-promotion evidence report
| +-- restore_model_artifacts.py # Restore model artifacts from a promotion backup
| +-- summarize_prediction_logs.py # Privacy-safe inference monitoring summary
| +-- package_retrain_artifacts.py # Zip/checksum retrain handoff bundle
| +-- prepare_gpu_retrain_handoff.py # Package reviewed safe data for Colab/GPU retrain
| +-- run_retrain_pipeline.py # End-to-end retrain validation/package runner
| +-- setup_retrain_environment.py # Install retrain dependencies before environment validation
| +-- validate_retrain_environment.py # Check packages/GPU before retraining
| +-- verify_retrain_bundle.py # Verify/downloaded retrain zip before promotion
+-- .github/workflows/
| +-- ml-quality-gate.yml # CI tests and artifact checks
+-- notebooks/
| +-- HCPG_GNN_Colab.ipynb # Exploratory notebook
+-- deployment/
| +-- Dockerfile
| +-- cloudbuild.yaml
| +-- app.yaml
+-- copy_of_vl_detection.py # Early prototype (parsing exploration)
+-- requirements.txt
```
## 快速开始
### 1. 本地开发
```
cd "final year project"
pip install -r requirements.txt
python -m uvicorn backend.app:app --reload
# 在浏览器中打开 frontend/index.html
```
Artifact 验证：
```
python scripts/validate_artifacts.py
```
### 2. Colab 训练（推荐）
1. 打开 [Google Colab](https://colab.research.google.com)，创建一个新的 notebook
2. 上传完整的 `models/` 目录或在 Colab 中克隆此仓库
3. 启用 GPU：运行时 > 更改运行时类型 > T4 GPU
4. 运行：
%cd /content/HCPG_GNN_Vulnerability_Detection-main
!PYTHONPATH="$PWD" python models/colab_hgt_visual.py
可选的额外标注数据集：
!EXTRA_DATASET_DIRS="/content/my_labelled_dataset" PYTHONPATH="$PWD" python models/colab_hgt_visual.py
可选的经过审查的安全合约：
!SAFE_DATASET_DIRS="/content/verified_safe_contracts" PYTHONPATH="$PWD" python models/colab_hgt_visual.py
从经过审查的候选合约中创建该安全文件夹：
!python scripts/build_safe_review_queue.py queue --candidate-dir /content/safe_contract_candidates --output /content/safe_review_queue.json
!python scripts/build_safe_review_queue.py export-review-sheet /content/safe_review_queue.json --output /content/safe_review_sheet.csv
# 手动审查/编辑 /content/safe_review_sheet.csv:
# 设置 status=approved_safe, reviewer, reviewed_at_utc, safe_for_supported_labels=true,
# 并将所有清单字段设置为 true，然后：
!python scripts/build_safe_review_queue.py apply-review-sheet /content/safe_review_queue.json --review-sheet /content/safe_review_sheet.csv --output /content/safe_review_queue_reviewed.json
!python scripts/build_safe_review_queue.py export-safe /content/safe_review_queue_reviewed.json --output-dir /content/verified_safe_contracts --label-manifest /content/safe_labels.json --min-items 100
该队列包含 `static_review_hints`，并且经过排序，以便优先审查具有低级调用、delegatecall、时间戳依赖、内联汇编或高复杂度的合约。这些提示仅作为分类辅助；它们永远不会自动将合约标记为安全。
可选的基于清单的多标签数据集：
!LABEL_MANIFESTS="/content/labels.json" PYTHONPATH="$PWD" python models/colab_hgt_visual.py
示例 `labels.json`：
{
"contracts": [
{"path": "contracts/Vault.sol", "labels": ["reentrancy", "SWC-115"]},
{"path": "contracts/SafeToken.sol", "safe": true}
]
}
在训练前验证清单：
!python scripts/validate_label_manifest.py /content/labels.json --require-safe --require-multilabel
在花费 GPU 时间之前，验证完整的本地重训练数据包：
!python scripts/setup_retrain_environment.py --report /content/retrain_environment_setup_report.json
!python scripts/validate_retrain_environment.py --require-cuda --report /content/retrain_environment_report.json
!python scripts/retrain_preflight.py --safe-dir /content/verified_safe_contracts --extra-dir /content/my_labelled_dataset --label-manifest /content/labels.json --report /content/retrain_preflight_report.json
要将已审查的安全文件夹打包为可重复的 Colab 上传 bundle：
python scripts/prepare_gpu_retrain_handoff.py \
--safe-dir /path/to/reviewed_safe_contracts \
--output-dir /path/to/gpu_retrain_handoff \
--min-safe-contracts 100 \
--force
上传生成的 `verified_safe_contracts.zip` 并按照生成的 `COLAB_RETRAIN_HANDOFF.ipynb` 或 `COLAB_RETRAIN_HANDOFF.md` 进行操作。
或者在一个命令中运行完整的重训练交接 pipeline：
!python scripts/run_retrain_pipeline.py \
--safe-dir /content/verified_safe_contracts \
--extra-dir /content/my_labelled_dataset \
--label-manifest /content/labels.json \
--min-safe-contracts 100
5. 下载 `model_metrics.json`、`model_dataset_manifest.json`、`model_error_analysis.json`、`model_split_manifest.json`、`baseline_comparison.json`、`calibration_report.json`、`calibration_curves.png`、`training_run_manifest.json`、`training_experiments.jsonl`、`best_hgt_model.pt` 和 `training_curves.png`
该脚本会自动克隆 SmartBugs 和 SolidiFI；无需手动下载数据集。新导出的 checkpoint 包含 API 在启用实时 GNN 推理之前进行验证的元数据，而旧版原始 state-dict checkpoint 仍然可以加载。在验证集 BCE 上拟合标量温度，然后在校准后的验证集划分上调整各类别的决策阈值，并由 `/api/model-info` 和 `/api/analyze` 提供。`LABEL_MANIFESTS` 支持对具有多个受支持问题的合约进行真正的多标签注释。每次训练运行后审查 `model_error_analysis.json`，以决定下一步的改进应该是更多的标注数据、更干净的标签、更安全的阈值，还是针对特定漏洞类别的特征工程。审查 `baseline_comparison.json` 以确认 GNN 在相同的留出划分上优于确定性特征基线，并在将概率视为置信度估计之前审查 `calibration_report.json` 和 `calibration_curves.png`。
在替换本地 artifact 之前，请运行：
```
python scripts/validate_retrain_outputs.py /path/to/colab/output --baseline-metrics models/model_metrics.json --min-safe-contracts 100
```
为同一候选人生成一份人类可读的审查报告：
```
python scripts/generate_model_quality_report.py /path/to/colab/output --baseline-metrics models/model_metrics.json --output /path/to/colab/output/model_quality_report.md
```
将经过验证的候选人打包为一个带有校验和的交接 zip 文件：
```
python scripts/package_retrain_artifacts.py /path/to/colab/output --baseline-metrics models/model_metrics.json --output /path/to/colab/output/model_artifact_bundle.zip
```
将该 zip 文件下载到本地后，在进行发布（promotion）之前验证校验和并将其解压：
```
python scripts/verify_retrain_bundle.py /path/to/model_artifact_bundle.zip \
--output-dir /path/to/verified_candidate \
--baseline-metrics models/model_metrics.json \
--min-safe-contracts 100
```
为了更严格地保护少数类，请添加各类别的下限：
```
python scripts/validate_retrain_outputs.py /path/to/colab/output --baseline-metrics models/model_metrics.json --min-safe-contracts 100 --min-class-f1 0.05 --min-class-recall 0.05
```
在确切的测试集划分上可选的静态分析器基线：
```
python scripts/evaluate_static_baseline.py --split-manifest /path/to/colab/output/model_split_manifest.json --run-slither --output /path/to/colab/output/static_baseline_comparison.json
```
当该静态基线可用时，在候选人验证期间强制执行它：
```
python scripts/validate_retrain_outputs.py /path/to/colab/output --baseline-metrics models/model_metrics.json --min-safe-contracts 100 --require-static-baseline
```
然后发布经过验证的候选人：
```
python scripts/promote_retrain_outputs.py /path/to/verified_candidate --baseline-metrics models/model_metrics.json --min-safe-contracts 100
```
发布后，在部署前运行行为金丝雀测试：
```
python scripts/run_model_canary_suite.py --report models/model_canary_report.json
```
或者在一个命令中运行完整的发布后验证堆栈：
```
python scripts/run_post_promotion_verification.py
```
如果必须回滚已发布的模型，请恢复由 `promote_retrain_outputs.py` 创建的备份：
```
python scripts/restore_model_artifacts.py models/artifact_backups/
python scripts/run_post_promotion_verification.py
```
然后运行最终的 ML 发布门禁：
```
python scripts/run_ml_release_gate.py --require-promotion-report --require-quality-report
```
可选的生产级预测监控：
```
$env:HCPG_ANALYSIS_AUDIT_LOG_PATH="logs/analysis_audit.jsonl"
python -m uvicorn backend.app:app --host 127.0.0.1 --port 8000
python scripts/summarize_prediction_logs.py logs/analysis_audit.jsonl \
--report logs/analysis_audit_summary.json \
--max-unsafe-rate 0.8 \
--max-p95-latency-ms 5000 \
--min-gnn-loaded-rate 0.9 \
--max-artifact-warning-rate 0.0
```
审计日志存储哈希值、标签、模型/运行时状态、风险摘要和计时。它不存储 Solidity 源代码、图 payload 或函数体。
`/api/model-info` 和 `/api/benchmark` 还公开了 `artifacts.validation`，这是对当前权重、指标文件、阈值、数据集来源计数、安全合约覆盖率、checkpoint 格式和训练后报告的结构化通过/警告/失败审计。警告意味着应用程序仍然可以运行，但在声明拥有完全最新的模型之前，应该重新生成或强化该 artifact。
如需直接进行机器可读的 artifact 审计，请调用 `/api/artifact-validation`。有关人类项目审计和后续 pipeline，请参阅 `docs/ML_ENGINEERING_AUDIT.md`。
### 3. Docker 部署
```
docker build -t hcpgnn-auditor .
docker run -p 8000:8000 hcpgnn-auditor
```
### 4. Google Cloud 部署
```
gcloud builds submit --config=deployment/cloudbuild.yaml
# 或者
gcloud run deploy hcpgnn-auditor \
--source . \
--region us-central1 \
--allow-unauthenticated
```
完整步骤请参见 `deployment/GOOGLE_CLOUD_DEPLOY.md`。
## API 用法
### 分析合约
```
curl -X POST "http://localhost:8000/analyze" \
-H "Content-Type: application/json" \
-d '{"source_code": "pragma solidity ^0.8.0; contract Test { }"}'
```
### 获取示例合约
```
curl "http://localhost:8000/sample/reentrancy"
```
## 漏洞类别
目前的行业重点遵循 OWASP 智能合约 Top 10 的侧重点：在 UI 中重点突出了 Access Control、Price Oracle Manipulation 和 Reentrancy。Access Control 和 Reentrancy 是经过训练的类别外加 regex 检查；由于当前数据集没有提供经过训练的预言机类别，Price Oracle Manipulation 仅进行结构化 regex 覆盖。
| SWC ID | 漏洞 | 严重性 | 数据集覆盖范围 |
|--------|--------------|----------|-----------------|
| SWC-107 | Reentrancy | 严重 | SmartBugs + SolidiFI |
|C-115 | Access Control | 严重 | SmartBugs + SolidiFI |
| SC02 | Price Oracle Manipulation | 严重/高 | 仅限 Regex 结构化检查 |
| SWC-101 | Integer Overflow/Underflow | 高 | SmartBugs + SolidiFI |
| SWC-114 | Transaction Ordering Dependency | 高 | SmartBugs + SolidiFI |
| SWC-104 | Unchecked Call Return Value | 中 | SmartBugs + SolidiFI |
| SWC-112 | Delegatecall to Untrusted Callee | 高 | SolidiFI |
## 环境要求
- Python 3.8+
- PyTorch 2.0+
- PyTorch Geometric
- FastAPI
- NetworkX
- scikit-learn
## 许可证
MIT License
## 参考文献
- Atzei et al. (2017). A survey of attacks on Ethereum smart contracts. *POST 2017*.
- Yamaguchi et al. (2014). Modeling and discovering vulnerabilities with code property graphs. *IEEE S&P*.
- Hu et al. (2020). Heterogeneous graph transformer. *WWW 2020*.
- Ferreira Torres et al. (2020). SmartBugs: A framework to analyze Solidity smart contracts. *ASE 2020*.
- Ghaleb & Pattabiraman (2020). How effective are smart contract analysis tools? *ISSTA 2020* (SolidiFI).
标签：AV绕过, FastAPI, Solidity, 云安全监控, 代码安全分析, 凭据扫描, 图神经网络, 智能合约审计, 网络测绘, 请求拦截, 逆向工具, 静态分析