Python脚本实现依赖漏洞自动扫描工具

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
🔧 Python 依赖漏洞自动扫描脚本
支持：requirements.txt / pyproject.toml
使用 Snyk API 查询已知漏洞（CVE）
"""

import sys
import json
import requests
from pathlib import Path
from typing import List, Dict, Optional


class VulnerabilityScanner:
    def __init__(self, api_key: str = None):
        self.api_key = api_key or "SNYK-KEY-PLACEHOLDER"
        self.base_url = "https://snyk.io/api/v1"
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"token {self.api_key}",
            "Content-Type": "application/json",
            "User-Agent": "OWASP-Scanner/1.0"
        })

    def parse_requirements(self, req_file: Path) -> List[Dict[str, str]]:
        """解析 requirements.txt 文件"""
        packages = []
        try:
            with open(req_file, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith('#'):
                        continue
                    # 支持格式：package==1.2.3, package>=1.0, package
                    parts = line.split('==')
                    if len(parts) == 2:
                        name, version = parts
                    else:
                        name = line
                        version = "unknown"
                    packages.append({"name": name.strip(), "version": version.strip()})
        except Exception as e:
            print(f"❌ 解析 {req_file} 失败: {e}")
            return []
        return packages

    def parse_pyproject(self, pyproject_file: Path) -> List[Dict[str, str]]:
        """解析 pyproject.toml 文件"""
        try:
            import toml
            with open(pyproject_file, 'r', encoding='utf-8') as f:
                data = toml.load(f)
            packages = []
            # 从 [build-system] 或 [project.dependencies] 中提取
            deps = data.get("project", {}).get("dependencies", [])
            for dep in deps:
                if "==" in dep:
                    name, version = dep.split("==", 1)
                else:
                    name = dep
                    version = "unknown"
                packages.append({"name": name.strip(), "version": version.strip()})
            return packages
        except Exception as e:
            print(f"❌ 解析 {pyproject_file} 失败: {e}")
            return []

    def get_vulnerabilities(self, package_name: str, version: str) -> List[Dict]:
        """查询 Snyk API 获取漏洞信息"""
        url = f"{self.base_url}/package/pypi/{package_name}/versions"
        try:
            response = self.session.get(url, timeout=10)
            if response.status_code == 404:
                return []
            if response.status_code != 200:
                print(f"⚠️  请求失败: {response.status_code} - {response.text}")
                return []

            data = response.json()
            vulns = []

            # 遍历所有版本，检查目标版本是否有漏洞
            for version_info in data.get("versions", []):
                if version_info["version"] == version:
                    for vuln in version_info.get("vulnerabilities", []):
                        vulns.append({
                            "id": vuln.get("id"),
                            "title": vuln.get("title"),
                            "severity": vuln.get("severity"),
                            "cvss": vuln.get("cvss", {}).get("baseScore"),
                            "published": vuln.get("published"),
                            "url": vuln.get("url")
                        })
                    break

            return vulns

        except Exception as e:
            print(f"❌ 查询 {package_name} {version} 时出错: {e}")
            return []

    def scan_project(self, project_root: Path):
        """主扫描函数"""
        print(f"🔍 正在扫描项目: {project_root.resolve()}")
        found_vulnerabilities = []

        # 尝试解析 requirements.txt
        req_file = project_root / "requirements.txt"
        if req_file.exists():
            print("📦 正在解析 requirements.txt...")
            packages = self.parse_requirements(req_file)
        else:
            packages = []

        # 如果没有 requirements.txt，尝试解析 pyproject.toml
        if not packages:
            pyproject_file = project_root / "pyproject.toml"
            if pyproject_file.exists():
                print("📦 正在解析 pyproject.toml...")
                packages = self.parse_pyproject(pyproject_file)
            else:
                print("❌ 未找到 requirements.txt 或 pyproject.toml")
                return

        # 扫描每个依赖
        for pkg in packages:
            name = pkg["name"]
            version = pkg["version"]
            print(f"🔍 正在检查: {name} ({version})")

            vulns = self.get_vulnerabilities(name, version)
            if vulns:
                print(f"🚨 发现漏洞: {name} ({version})")
                for v in vulns:
                    print(f"   • {v['id']} | {v['title']} | {v['severity']} | CVSS: {v['cvss']}")
                found_vulnerabilities.extend(vulns)
            else:
                print(f"✅ 安全: {name} ({version}) 无已知漏洞")

        # 输出总结
        if found_vulnerabilities:
            print("\n" + "="*60)
            print("🚨 **发现以下已知漏洞**")
            print("="*60)
            for v in found_vulnerabilities:
                print(f"- {v['id']}: {v['title']} | {v['severity']} | {v['url']}")
            print("="*60)
            print("💡 建议：升级依赖版本或使用 Snyk 修复命令")
        else:
            print("\n🎉 ✅ 所有依赖项均无已知漏洞！")

        return len(found_vulnerabilities) > 0


def main():
    import argparse

    parser = argparse.ArgumentParser(description="自动扫描 Python 项目中的已知漏洞依赖")
    parser.add_argument("project_root", nargs="?", default=".", help="项目根目录（默认为当前目录）")
    parser.add_argument("--api-key", help="Snyk API Key（可选，若不提供则使用模拟数据）")
    args = parser.parse_args()

    project_root = Path(args.project_root).resolve()

    if not project_root.exists():
        print(f"❌ 项目目录不存在: {project_root}")
        sys.exit(1)

    scanner = VulnerabilityScanner(api_key=args.api_key)
    scanner.scan_project(project_root)


if __name__ == "__main__":
    main()

# 1. 保存脚本为 scan_python_vulnerabilities.py
# 2. 安装依赖（仅需 requests）
pip install requests

# 3. 运行扫描（推荐使用 API Key）
python scan_python_vulnerabilities.py . --api-key "SNYK-KEY-PLACEHOLDER"

# 或扫描指定项目
python scan_python_vulnerabilities.py /path/to/your/project --api-key "your-snyk-key"