#!/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 = []

            # 遍歷所有版本，檢查目標(biāo)版本是否有漏洞
            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):
        """主掃描函數(shù)"""
        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"?? 發(fā)現(xiàn)漏洞: {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}) 無已知漏洞")

        # 輸出總結(jié)
        if found_vulnerabilities:
            print("\n" + "="*60)
            print("?? **發(fā)現(xiàn)以下已知漏洞**")
            print("="*60)
            for v in found_vulnerabilities:
                print(f"- {v['id']}: {v['title']} | {v['severity']} | {v['url']}")
            print("="*60)
            print("?? 建議：升級依賴版本或使用 Snyk 修復(fù)命令")
        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="項目根目錄（默認為當(dāng)前目錄）")
    parser.add_argument("--api-key", help="Snyk API Key（可選，若不提供則使用模擬數(shù)據(jù)）")
    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"