ai-audit-tool 0.2.0

AI包漏洞审计工具，支持命令行和模块调用

AI Package Auditor: 智能包漏洞审计工具

AI Package Auditor 是一款专为AI/机器学习项目设计的包漏洞审计工具，支持命令行独立运行和代码模块调用两种模式。通过静态分析与AI增强审计结合，帮助开发者识别依赖包中的安全风险，特别是针对AI框架（如TensorFlow、PyTorch）的特定漏洞。

核心功能

基础审计能力

• 依赖自动收集：扫描项目依赖文件（requirements.txt、pyproject.toml等）及已安装包
• AI包智能分类：按领域（自然语言处理、计算机视觉等）识别AI相关包
• 多源漏洞查询：整合OSV等数据库，检测已知安全漏洞
• 风险等级评估：基于漏洞 severity、影响范围等因素评分（0-10分）
• 多格式报告：支持控制台输出和JSON格式报告，包含修复建议

AI增强功能（可选）

当配置AI模型后，工具将额外提供：

• 代码级安全审计：对核心AI代码进行深度分析
• 上下文感知检测：识别AI特有的风险（如模型加载漏洞、数据处理安全）
• 智能修复建议：针对AI框架特性提供定制化修复方案

安装指南

前置要求

• Python 3.8+
• 网络连接（用于漏洞数据库查询和AI API调用）

安装方式

1. 从Wheel包安装（推荐）# 下载或构建whl包后安装

pip install ai_audit_tool-0.1.0-py3-none-any.whl

如需启用AI功能，安装可选依赖

pip install ai_audit_tool-0.1.0-py3-none-any.whl[ai]

2. 源码安装# 克隆仓库

git clone https://github.com/honysyang/ai-package-auditor.git cd ai-package-auditor

安装

pip install .

安装开发依赖（用于二次开发）

pip install .[dev]

使用指南

模式1：命令行工具

安装后可直接通过ai-audit命令运行：

基础审计（当前目录，控制台输出报告）

ai-audit

自定义项目路径并生成JSON报告

ai-audit --project-path ./my_ai_project --json --json-path audit_report.json

启用AI审计（通过命令行参数配置AI模型）

ai-audit
--project-path ./my_ai_project
--ai-api-base "https://api.chatanywhere.tech/v1/chat/completions"
--ai-api-key "your-api-key"
--ai-model-name "gpt-4o-mini-ca"
--max-workers 8
--json

命令行参数说明

参数	描述	默认值
--project-path	目标项目路径	.（当前目录）
--max-workers	并行查询线程数	5
--print/--no-print	是否在控制台打印报告	启用（--print）
--json	是否生成JSON报告	禁用
--json-path	JSON报告路径	ai_audit_report.json
--ai-api-base	AI模型API地址（启用AI功能需提供）	-
--ai-api-key	AI模型API密钥（启用AI功能需提供）	-
--ai-model-name	AI模型名称（启用AI功能需提供）	-

模式2：代码模块调用

作为Python模块集成到现有工作流中： from ai_audit import AIPackageAuditor, AIModelConfig

1. 基础审计（无AI功能）

auditor = AIPackageAuditor(max_workers=5) results = auditor.audit( project_path="./my_ai_project", report_options={ "print": True, "json": True, "json_path": "custom_report.json" } )

2. 启用AI审计

ai_config = AIModelConfig( api_base="https://api.chatanywhere.tech/v1/chat/completions", api_key="your-api-key", model_name="gpt-4o-mini-ca", timeout=30 # 超时时间（秒） )

auditor = AIPackageAuditor( max_workers=8, ai_config=ai_config )

results = auditor.audit( project_path="./my_ai_project", report_options={"print": True, "json": True} )

解析审计结果

high_risk_pkgs = results["base_audit"]["summary"]["high_risk_packages"] print(f"发现 {len(high_risk_pkgs)} 个高风险包")

AI模型配置

AI功能支持三种配置方式（优先级从高到低）：

1. 配置文件（推荐）

创建ai_config.json：{ "api_base": "https://api.chatanywhere.tech/v1/chat/completions", "api_key": "your-api-key", "model_name": "gpt-4o-mini-ca", "timeout": 30, "max_tokens": 2048 }使用时指定配置文件：ai-audit --ai-config ai_config.json

2. 环境变量export AI_API_BASE="https://api.chatanywhere.tech/v1/chat/completions"

export AI_API_KEY="your-api-key" export AI_MODEL_NAME="gpt-4o-mini-ca" ai-audit

3. 命令行参数

如「使用指南」中命令行示例所示。

报告解读

工具生成两种报告（取决于是否启用AI功能）：

1. 基础审计报告

包含依赖分析和漏洞评估：{ "meta": { "audit_time": "2023-10-01 15:30:00", "project_path": "/path/to/project", "total_dependencies": 28, "total_ai_packages": 5 }, "summary": { "high_risk_count": 2, "medium_risk_count": 3, "low_risk_count": 1 }, "vulnerable_packages": [ { "package": "tensorflow", "version": "2.5.0", "risk_score": 8.5, "vulnerabilities": [ { "id": "CVE-2023-1234", "description": "TensorFlow存在代码执行漏洞...", "fix_versions": ["2.10.0"] } ], "mitigation": "建议升级到2.10.0版本..." } ] }

2. AI代码审计报告（启用AI时）

额外包含代码级分析：{ "ai_audit_results": [ { "file_path": "models/trainer.py", "risk_score": 7.2, "issues": [ { "description": "模型加载未验证文件完整性，可能导致恶意模型执行", "suggestion": "添加模型哈希校验，使用tf.saved_model.load时验证签名" } ] } ] }

工程化设计

项目结构ai_audit/

├── init.py # 包导出声明 ├── auditor.py # 审计主逻辑 ├── cli.py # 命令行入口 ├── models/ # 数据模型（依赖、漏洞、AI配置等） ├── collectors/ # 依赖收集模块 ├── assessors/ # 风险评估模块 ├── services/ # 漏洞查询服务 └── ai_auditor.py # AI审计模块（条件加载）

核心技术特点

• 条件性功能加载：AI模块仅在配置完整时加载，不影响基础功能
• 多源配置兼容：支持文件、环境变量、命令行参数混合配置
• 并行化查询：通过线程池加速漏洞数据库查询
• 跨平台兼容：路径处理适配Windows/macOS/Linux
• 可扩展架构：模块化设计便于添加新的漏洞数据源或AI模型

开发指南

本地开发# 安装开发依赖

pip install .[dev]

运行测试

pytest tests/

代码格式化

black ai_audit/

构建Wheel包

python -m build

贡献指南

1. Fork 仓库并创建分支（feature/xxx 或 fix/xxx）
2. 提交代码前确保测试通过
3. 提交PR时请说明功能变更或修复内容

许可证

本项目基于 MIT许可证 开源，允许自由使用、修改和分发。