|
| 1 | +"""用于管理每个案例分析输入的工作区助手。""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import logging |
| 6 | +import re |
| 7 | +from dataclasses import dataclass |
| 8 | +from pathlib import Path |
| 9 | +from typing import Dict, Iterable, Optional |
| 10 | + |
| 11 | +from .cve_fetcher import fetch_cve_from_nvd, save_cve_data |
| 12 | + |
| 13 | +CVE_FILE_PATTERN = re.compile(r"(CVE-\d{4}-\d+)", re.IGNORECASE) |
| 14 | +logger = logging.getLogger(__name__) |
| 15 | + |
| 16 | + |
| 17 | +@dataclass(frozen=True) |
| 18 | +class CasePaths: |
| 19 | + """单个案例分析已解析的位置。""" |
| 20 | + |
| 21 | + root: Path |
| 22 | + source_code: Path |
| 23 | + db: Path |
| 24 | + inputs: Path |
| 25 | + intel: Path |
| 26 | + |
| 27 | + |
| 28 | +@dataclass(frozen=True) |
| 29 | +class ExtraFile: |
| 30 | + """额外输入文件的元数据。""" |
| 31 | + |
| 32 | + path: Path |
| 33 | + |
| 34 | + def read_text(self) -> str: |
| 35 | + """读取文件内容为文本。""" |
| 36 | + return self.path.read_text(encoding='utf-8') |
| 37 | + |
| 38 | + |
| 39 | +@dataclass(frozen=True) |
| 40 | +class CveAssets: |
| 41 | + """案例中单个CVE的本地工件。""" |
| 42 | + |
| 43 | + cve_id: str |
| 44 | + json_path: Path |
| 45 | + diff_path: Optional[Path] |
| 46 | + extra_files: tuple[ExtraFile, ...] = () # 额外输入文件 |
| 47 | + |
| 48 | + def has_extra_files(self) -> bool: |
| 49 | + """检查是否有额外文件。""" |
| 50 | + return len(self.extra_files) > 0 |
| 51 | + |
| 52 | + def get_all_extra_content(self) -> str: |
| 53 | + """获取所有额外文件的内容(用于分析上下文)。""" |
| 54 | + if not self.extra_files: |
| 55 | + return "" |
| 56 | + |
| 57 | + sections = [f"=== 额外输入文件 ({len(self.extra_files)} 个) ===\n"] |
| 58 | + |
| 59 | + for extra_file in self.extra_files: |
| 60 | + sections.append(f"\n--- 文件: {extra_file.path.name} ---") |
| 61 | + try: |
| 62 | + content = extra_file.read_text() |
| 63 | + sections.append(content) |
| 64 | + except Exception as e: |
| 65 | + sections.append(f"[无法读取文件: {e}]") |
| 66 | + |
| 67 | + return "\n".join(sections) |
| 68 | + |
| 69 | + |
| 70 | +def resolve_case(case_id: str, *, base_dir: Path = Path("projects")) -> CasePaths: |
| 71 | + """解析并验证案例工作区。""" |
| 72 | + |
| 73 | + if not case_id or Path(case_id).is_absolute(): |
| 74 | + raise ValueError(f"Invalid case id: {case_id!r}") |
| 75 | + |
| 76 | + base_root = base_dir.resolve() |
| 77 | + root = (base_root / case_id).resolve() |
| 78 | + if not root.is_relative_to(base_root): |
| 79 | + raise ValueError(f"Case path escapes projects directory: {case_id!r}") |
| 80 | + |
| 81 | + mapping = { |
| 82 | + "source_code": root / "source_code", |
| 83 | + "db": root / "db", |
| 84 | + "inputs": root / "inputs", |
| 85 | + "intel": root / "intel", |
| 86 | + } |
| 87 | + |
| 88 | + missing = [] |
| 89 | + not_dirs = [] |
| 90 | + for name, path in mapping.items(): |
| 91 | + if not path.exists(): |
| 92 | + missing.append(name) |
| 93 | + elif not path.is_dir(): |
| 94 | + not_dirs.append(f"{name} (路径存在但不是目录: {path})") |
| 95 | + |
| 96 | + if missing: |
| 97 | + raise FileNotFoundError( |
| 98 | + f"Case '{case_id}' is missing required directories: {', '.join(missing)}" |
| 99 | + ) |
| 100 | + |
| 101 | + if not_dirs: |
| 102 | + raise ValueError( |
| 103 | + f"Case '{case_id}' has invalid paths (expected directories): {', '.join(not_dirs)}" |
| 104 | + ) |
| 105 | + |
| 106 | + return CasePaths( |
| 107 | + root=root, |
| 108 | + source_code=mapping["source_code"], |
| 109 | + db=mapping["db"], |
| 110 | + inputs=mapping["inputs"], |
| 111 | + intel=mapping["intel"], |
| 112 | + ) |
| 113 | + |
| 114 | + |
| 115 | +def _discover_extra_files(inputs_dir: Path, cve_id: str) -> tuple[ExtraFile, ...]: |
| 116 | + """ |
| 117 | + 发现 inputs 目录中的额外文件。 |
| 118 | + |
| 119 | + 排除标准的 CVE JSON、diff 和 patch 文件。 |
| 120 | + """ |
| 121 | + extra_files = [] |
| 122 | + |
| 123 | + for file_path in inputs_dir.iterdir(): |
| 124 | + if not file_path.is_file(): |
| 125 | + continue |
| 126 | + |
| 127 | + # 排除标准 CVE 文件 |
| 128 | + if file_path.name.startswith('CVE-') and file_path.suffix in ('.json', '.diff', '.patch'): |
| 129 | + continue |
| 130 | + |
| 131 | + # 排除隐藏文件和临时文件 |
| 132 | + if file_path.name.startswith('.') or file_path.name.endswith('~'): |
| 133 | + continue |
| 134 | + |
| 135 | + extra_file = ExtraFile(path=file_path) |
| 136 | + extra_files.append(extra_file) |
| 137 | + |
| 138 | + # 按文件名排序 |
| 139 | + extra_files.sort(key=lambda f: f.path.name) |
| 140 | + |
| 141 | + if extra_files: |
| 142 | + logger.info(f"📂 [额外文件] 发现 {len(extra_files)} 个额外输入文件:") |
| 143 | + for extra_file in extra_files: |
| 144 | + logger.info(f" - {extra_file.path.name}") |
| 145 | + |
| 146 | + return tuple(extra_files) |
| 147 | + |
| 148 | + |
| 149 | +def discover_cve_assets(case_paths: CasePaths) -> CveAssets: |
| 150 | + """ |
| 151 | + 在案例输入目录中定位CVE JSON/diff对,支持文件缺失时的回退机制。 |
| 152 | +
|
| 153 | + 自动选择第一个可用的CVE进行,如果JSON文件缺失则从NVD API获取。 |
| 154 | + 同时发现并分类额外的输入文件。 |
| 155 | + """ |
| 156 | + |
| 157 | + json_files = sorted(case_paths.inputs.glob("CVE-*.json")) |
| 158 | + json_map: Dict[str, Path] = {} |
| 159 | + cve_id: str |
| 160 | + json_path: Path |
| 161 | + |
| 162 | + # 处理本地JSON文件 |
| 163 | + if json_files: |
| 164 | + for path in json_files: |
| 165 | + cve_id_extracted = extract_cve_id(path.name) |
| 166 | + if not cve_id_extracted: |
| 167 | + continue |
| 168 | + json_map[cve_id_extracted.upper()] = path |
| 169 | + |
| 170 | + if json_map: |
| 171 | + # 使用本地找到的JSON文件 |
| 172 | + cve_id = sorted(json_map.keys())[0] |
| 173 | + json_path = json_map[cve_id] |
| 174 | + logger.info(f"📁 [本地文件] 找到CVE JSON文件: {json_path}") |
| 175 | + else: |
| 176 | + # 本地有JSON文件但格式都无效,尝试从diff/patch文件推断CVE ID |
| 177 | + diff_files = sorted(case_paths.inputs.glob("CVE-*.diff")) |
| 178 | + patch_files = sorted(case_paths.inputs.glob("CVE-*.patch")) |
| 179 | + all_diff_patch_files = sorted(diff_files + patch_files) |
| 180 | + if all_diff_patch_files: |
| 181 | + cve_id = extract_cve_id(all_diff_patch_files[0].name) |
| 182 | + if cve_id: |
| 183 | + logger.info(f"📝 [推断ID] 从diff/patch文件推断CVE ID: {cve_id}") |
| 184 | + else: |
| 185 | + raise ValueError( |
| 186 | + f"Inputs directory {case_paths.inputs} contains files but no valid CVE IDs could be extracted" |
| 187 | + ) |
| 188 | + else: |
| 189 | + raise ValueError( |
| 190 | + f"Inputs directory {case_paths.inputs} contains JSON files but no valid CVE IDs could be extracted" |
| 191 | + ) |
| 192 | + else: |
| 193 | + # 没有本地JSON文件,尝试从diff/patch文件推断CVE ID |
| 194 | + diff_files = sorted(case_paths.inputs.glob("CVE-*.diff")) |
| 195 | + patch_files = sorted(case_paths.inputs.glob("CVE-*.patch")) |
| 196 | + all_diff_patch_files = sorted(diff_files + patch_files) |
| 197 | + if all_diff_patch_files: |
| 198 | + cve_id = extract_cve_id(all_diff_patch_files[0].name) |
| 199 | + if not cve_id: |
| 200 | + raise ValueError( |
| 201 | + f"Inputs directory {case_paths.inputs} contains diff/patch files but no valid CVE IDs could be extracted" |
| 202 | + ) |
| 203 | + logger.info(f"📝 [推断ID] 未找到JSON文件,从diff/patch文件推断CVE ID: {cve_id}") |
| 204 | + else: |
| 205 | + raise FileNotFoundError( |
| 206 | + f"No CVE JSON, diff or patch files found in {case_paths.inputs}" |
| 207 | + ) |
| 208 | + |
| 209 | + # 如果没有有效的本地JSON文件,从NVD API获取 |
| 210 | + if not json_map: |
| 211 | + try: |
| 212 | + logger.info(f"🌐 [网络获取] 正在从NVD API获取CVE数据: {cve_id}") |
| 213 | + cve_data = fetch_cve_from_nvd(cve_id) |
| 214 | + json_path = save_cve_data(cve_id, cve_data, case_paths.inputs) |
| 215 | + logger.info(f"✅ [获取成功] CVE数据已保存到: {json_path}") |
| 216 | + except Exception as e: |
| 217 | + logger.error(f"❌ [获取失败] 无法获取CVE数据 {cve_id}: {e}") |
| 218 | + raise RuntimeError(f"Failed to fetch CVE data for {cve_id}: {e}") |
| 219 | + |
| 220 | + # 处理diff/patch文件(可选,缺失时不抛出异常) |
| 221 | + # 优先使用diff文件,如果没有则使用patch文件 |
| 222 | + diff_candidates = sorted(case_paths.inputs.glob("CVE-*.diff")) |
| 223 | + patch_candidates = sorted(case_paths.inputs.glob("CVE-*.patch")) |
| 224 | + all_candidates = diff_candidates + patch_candidates |
| 225 | + |
| 226 | + diff_map = {extract_cve_id(path.name): path for path in all_candidates} |
| 227 | + diff_path = diff_map.get(cve_id.upper()) |
| 228 | + |
| 229 | + if diff_path: |
| 230 | + file_type = "diff" if diff_path.suffix == ".diff" else "patch" |
| 231 | + logger.info(f"📄 [本地文件] 找到{file_type}文件: {diff_path}") |
| 232 | + else: |
| 233 | + logger.info(f"⚠️ [文件缺失] 未找到 {cve_id} 的diff/patch文件,将继续进行分析(无diff模式)") |
| 234 | + |
| 235 | + # 发现额外输入文件 |
| 236 | + extra_files = _discover_extra_files(case_paths.inputs, cve_id) |
| 237 | + |
| 238 | + return CveAssets( |
| 239 | + cve_id=cve_id, |
| 240 | + json_path=json_path, |
| 241 | + diff_path=diff_path, |
| 242 | + extra_files=extra_files |
| 243 | + ) |
| 244 | + |
| 245 | + |
| 246 | +def extract_cve_id(filename: str) -> Optional[str]: |
| 247 | + """从文件名中提取CVE标识符。""" |
| 248 | + |
| 249 | + match = CVE_FILE_PATTERN.search(filename) |
| 250 | + if match: |
| 251 | + return match.group(1).upper() |
| 252 | + return None |
| 253 | + |
| 254 | + |
| 255 | +def default_language_db(case_paths: CasePaths, language: str) -> Optional[Path]: |
| 256 | + """如果存在,返回给定语言的预期CodeQL数据库路径。""" |
| 257 | + |
| 258 | + language = language.lower() |
| 259 | + candidate = case_paths.db / language |
| 260 | + if candidate.exists(): |
| 261 | + return candidate |
| 262 | + if case_paths.db.exists() and any(case_paths.db.iterdir()): |
| 263 | + return case_paths.db |
| 264 | + return None |
0 commit comments