tools/exploit_tool.py

# tools/exploit_tool.py
# 功能：GitHub Exploit / PoC 搜尋 Tool
# Harness 支柱：Graceful Degradation（降級瀑布）+ Observability（原子化日誌）
# 擁有者：成員 C（Analyst Agent Pipeline）
#
# 使用方式：
#   from tools.exploit_tool import search_exploits
#
# 架構定位：
#   Analyst Agent 的「第二隻手」— 搜尋 CVE 的公開 Exploit/PoC
#   有公開 exploit = 攻擊門檻極低 = 風險指標 HIGH

import json
import os
import time
import hashlib
import logging
from datetime import datetime, timezone

import requests

logger = logging.getLogger("ThreatHunter")

# ══════════════════════════════════════════════════════════════
# 常數
# ══════════════════════════════════════════════════════════════

GITHUB_API_BASE = "https://api.github.com/search/repositories"
RESULTS_PER_PAGE = 10  # 限制結果數量，避免 context 過長
REQUEST_TIMEOUT = 20  # 秒

# Rate limit 控制
# GitHub 無認證 10 req/min，有認證 30 req/min
RATE_LIMIT_WITH_TOKEN = 2.0    # 有 Token: 30 req/min → 2s 間隔（保守）
RATE_LIMIT_WITHOUT_TOKEN = 6.0  # 無 Token: 10 req/min → 6s 間隔（保守）
MAX_RETRIES = 2

# 離線快取
CACHE_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
CACHE_TTL = 3600 * 24  # 24 小時過期

# 上次請求時間（模組級 rate limiter）
_last_request_time = 0.0


# ══════════════════════════════════════════════════════════════
# 輔助函式
# ══════════════════════════════════════════════════════════════

def _get_github_token() -> str:
    """取得 GitHub Token（優先從 config，備選環境變數）"""
    try:
        from config import GITHUB_TOKEN
        if GITHUB_TOKEN:
            return GITHUB_TOKEN
    except ImportError:
        pass
    return os.getenv("GITHUB_TOKEN", "")


def _get_cache_path(cve_id: str) -> str:
    """取得離線快取檔案路徑"""
    safe_name = hashlib.md5(cve_id.encode()).hexdigest()[:12]
    return os.path.join(CACHE_DIR, f"exploit_cache_{cve_id}_{safe_name}.json")


def _read_cache(cve_id: str) -> dict | None:
    """讀取離線快取，過期或不存在回傳 None"""
    cache_path = _get_cache_path(cve_id)
    try:
        if os.path.exists(cache_path):
            with open(cache_path, "r", encoding="utf-8") as f:
                cached = json.load(f)
            cached_time = cached.get("_cached_at", 0)
            if time.time() - cached_time < CACHE_TTL:
                logger.info("[OK] Exploit cache hit: %s", cve_id)
                return cached
            else:
                logger.info("[INFO] Exploit cache expired: %s", cve_id)
    except (json.JSONDecodeError, IOError) as e:
        logger.warning("[WARN] Exploit cache read failed: %s", e)
    return None


def _write_cache(cve_id: str, data: dict) -> None:
    """寫入離線快取"""
    try:
        os.makedirs(CACHE_DIR, exist_ok=True)
        cache_path = _get_cache_path(cve_id)
        data["_cached_at"] = time.time()
        with open(cache_path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    except (IOError, PermissionError) as e:
        logger.warning("[WARN] Exploit cache write failed: %s", e)


def _rate_limit() -> None:
    """Rate limiter — 確保不超過 GitHub API 限速"""
    global _last_request_time
    token = _get_github_token()
    interval = RATE_LIMIT_WITH_TOKEN if token else RATE_LIMIT_WITHOUT_TOKEN
    elapsed = time.time() - _last_request_time
    if elapsed < interval:
        wait = interval - elapsed
        logger.info("[WAIT] GitHub rate limit: waiting %.1fs", wait)
        time.sleep(wait)
    _last_request_time = time.time()


def _classify_exploit_type(name: str, description: str) -> str:
    """
    根據 repo 名稱和描述判斷 exploit 類型。

    分類規則：
      - 含 scanner/detect/check → scanner
      - 含 weaponized/exploit-db/metasploit → weaponized
      - 其他（poc/proof/test/demo 等）→ poc
    """
    text = f"{name} {description}".lower()

    # 掃描器 / 偵測工具
    scanner_keywords = ["scanner", "detect", "check", "scan", "checker", "finder"]
    for kw in scanner_keywords:
        if kw in text:
            return "scanner"

    # 武器化 exploit
    weaponized_keywords = ["weaponized", "exploit-db", "metasploit", "payload", "shellcode", "reverse shell"]
    for kw in weaponized_keywords:
        if kw in text:
            return "weaponized"

    # 預設為 PoC
    return "poc"


def _determine_risk_indicator(exploit_count: int, api_available: bool) -> str:
    """
    根據 exploit 數量和 API 可用性判定風險指標。

    邏輯：
      exploit_count > 0 → "HIGH"（有公開 exploit = 高風險）
      exploit_count == 0 且 API 可用 → "LOW"（確認沒有公開 exploit）
      API 不可用 → "UNKNOWN (API limited)"
    """
    if not api_available:
        return "UNKNOWN (API limited)"
    if exploit_count > 0:
        return "HIGH"
    return "LOW"


# ══════════════════════════════════════════════════════════════
# 核心查詢邏輯
# ══════════════════════════════════════════════════════════════

def _query_github_api(cve_id: str) -> dict | None:
    """
    呼叫 GitHub Search API，回傳原始 JSON response dict。
    失敗回傳 None。

    包含：
      - Rate limiting
      - 重試機制（最多 MAX_RETRIES 次）
      - Timeout 處理
    """
    token = _get_github_token()

    headers = {
        "Accept": "application/vnd.github.v3+json",
    }
    if token:
        headers["Authorization"] = f"token {token}"

    # 搜尋策略：用 CVE ID 作為關鍵字 + exploit 限定
    params = {
        "q": f"{cve_id} exploit",
        "sort": "stars",
        "order": "desc",
        "per_page": RESULTS_PER_PAGE,
    }

    for attempt in range(1, MAX_RETRIES + 1):
        _rate_limit()
        try:
            logger.info("[QUERY] GitHub API search: %s (attempt %d)", cve_id, attempt)
            response = requests.get(
                GITHUB_API_BASE,
                params=params,
                headers=headers,
                timeout=REQUEST_TIMEOUT,
            )

            if response.status_code == 200:
                return response.json()

            if response.status_code in (403, 429):
                # Rate limit 被觸發
                retry_after = response.headers.get("Retry-After", "60")
                logger.warning(
                    "[WARN] GitHub API %d (rate limited), Retry-After: %ss",
                    response.status_code, retry_after
                )
                time.sleep(min(int(retry_after), 30))  # 最多等 30 秒
                continue

            if response.status_code >= 500:
                logger.warning("[WARN] GitHub API %d (server error)", response.status_code)
                time.sleep(2)
                continue

            # 其他錯誤碼
            logger.warning("[WARN] GitHub API returned %d: %s", response.status_code, response.text[:200])
            return None

        except requests.exceptions.Timeout:
            logger.warning("[WARN] GitHub API timeout (%ds)", REQUEST_TIMEOUT)
            continue
        except requests.exceptions.ConnectionError:
            logger.warning("[WARN] GitHub API connection failed (network issue)")
            continue
        except requests.exceptions.RequestException as e:
            logger.warning("[WARN] GitHub API request error: %s", e)
            return None

    return None  # 所有重試都失敗


def _parse_github_response(raw: dict, cve_id: str) -> dict:
    """
    將 GitHub Search API 原始 response 轉換為 Tool 輸出格式。

    轉換 mapping：
      response.items[].full_name → repo_name
      response.items[].html_url → url
      response.items[].stargazers_count → stars
      response.items[].language → language
      response.items[].updated_at → last_updated
      response.items[].description → description
    """
    exploits = []
    items = raw.get("items", [])

    for repo in items:
        repo_name = repo.get("full_name", "")
        description = repo.get("description", "") or ""
        exploit_type = _classify_exploit_type(repo_name, description)

        # 提取最後更新日期（只取日期部分）
        updated_at = repo.get("updated_at", "")
        if "T" in updated_at:
            updated_at = updated_at.split("T")[0]

        exploits.append({
            "repo_name": repo_name,
            "url": repo.get("html_url", ""),
            "stars": repo.get("stargazers_count", 0),
            "language": repo.get("language", "") or "Unknown",
            "last_updated": updated_at,
            "description": description[:300],  # 截斷過長描述
            "type": exploit_type,
        })

    exploit_count = len(exploits)
    risk_indicator = _determine_risk_indicator(exploit_count, api_available=True)

    return {
        "cve_id": cve_id,
        "source": "GitHub API",
        "exploit_count": exploit_count,
        "exploits": exploits,
        "risk_indicator": risk_indicator,
    }


def _search_exploits_impl(cve_id: str) -> str:
    """
    search_exploits 的核心實作（與 CrewAI @tool 解耦，方便單元測試）。

    降級瀑布：
      1. 查 GitHub Search API → 成功則快取
      2. API 失敗（403/429 rate limit）→ 讀離線快取
      3. 快取也沒有 → 回傳 exploit_count: 0, risk_indicator: "UNKNOWN"
      4. 任何未預期錯誤 → 回傳安全的預設結果（絕不 crash）
    """
    try:
        # 清理 CVE ID（去空白、統一大寫）
        cve_id = cve_id.strip().upper()
        if not cve_id:
            logger.warning("[WARN] Exploit Tool received empty CVE ID input")
            return json.dumps({
                "cve_id": "",
                "source": "GitHub API (error)",
                "exploit_count": 0,
                "exploits": [],
                "risk_indicator": "UNKNOWN",
                "error": "No CVE ID provided",
            }, ensure_ascii=False, indent=2)

        logger.info("[QUERY] Exploit search: %s", cve_id)

        # 嘗試 GitHub API 查詢
        raw = _query_github_api(cve_id)

        if raw is not None:
            result = _parse_github_response(raw, cve_id)

            # 寫入快取供離線使用
            _write_cache(cve_id, result)

            logger.info(
                "[OK] Exploit search success: %s -> %d exploit(s), risk_indicator=%s",
                cve_id, result['exploit_count'], result['risk_indicator']
            )
            return json.dumps(result, ensure_ascii=False, indent=2)

        # API 失敗 → 嘗試讀快取
        cached = _read_cache(cve_id)
        if cached:
            cached.pop("_cached_at", None)
            cached["source"] = "GitHub API (cache)"
            cached["error"] = f"GitHub API unavailable, using cached data for '{cve_id}'"
            logger.info("[OK] Exploit using cache: %s", cve_id)
            return json.dumps(cached, ensure_ascii=False, indent=2)

        # 完全沒有資料
        empty_result = {
            "cve_id": cve_id,
            "source": "GitHub API (unavailable)",
            "exploit_count": 0,
            "exploits": [],
            "risk_indicator": "UNKNOWN (API limited)",
            "error": f"GitHub API unavailable and no cache for '{cve_id}'",
        }
        logger.info("[INFO] Exploit no data for: %s", cve_id)
        return json.dumps(empty_result, ensure_ascii=False, indent=2)

    except Exception as e:
        # 最後一道防線：任何未預期錯誤都不能讓 Agent crash
        logger.error("[FAIL] Exploit Tool unexpected error: %s", e, exc_info=True)
        error_result = {
            "cve_id": cve_id if 'cve_id' in dir() else "",
            "source": "GitHub API (error)",
            "exploit_count": 0,
            "exploits": [],
            "risk_indicator": "UNKNOWN",
            "error": f"Unexpected error: {str(e)}",
        }
        return json.dumps(error_result, ensure_ascii=False, indent=2)


# ══════════════════════════════════════════════════════════════
# CrewAI @tool 包裝（Agent 呼叫用）
# ══════════════════════════════════════════════════════════════

# ⚠️ 重要：使用「延遲載入」模式（LazyToolLoader）
# 原因：避免在 import 階段就觸發 CrewAI 的 tool 註冊
def _create_tool():
    """延遲建立 CrewAI Tool，僅在 Agent 實際使用時才 import"""
    from crewai.tools import tool

    @tool("search_exploits")
    def search_exploits(cve_id: str) -> str:
        """搜尋指定 CVE 的公開 Exploit 和 PoC 程式碼（透過 GitHub API）。
輸入單一 CVE ID（如 "CVE-2021-44228"），回傳公開 exploit 的數量、連結、星數等資訊。
有公開 exploit = 攻擊門檻極低 = 風險指標 HIGH。"""
        return _search_exploits_impl(cve_id)

    return search_exploits


# ── 延遲載入機制（與 nvd_tool.py 相同模式）──────────────────

class _LazyToolLoader:
    def __init__(self):
        self._tool = None

    def _load(self):
        if self._tool is None:
            self._tool = _create_tool()

    @property
    def search_exploits(self):
        self._load()
        return self._tool


_loader = _LazyToolLoader()


def __getattr__(name):
    """模組層級 __getattr__，支援 from tools.exploit_tool import search_exploits"""
    if name == "search_exploits":
        return _loader.search_exploits
    raise AttributeError(f"module 'tools.exploit_tool' has no attribute {name!r}")