diff options
Diffstat (limited to 'query.py')
| -rwxr-xr-x | query.py | 848 |
1 files changed, 848 insertions, 0 deletions
diff --git a/query.py b/query.py new file mode 100755 index 0000000..6680815 --- /dev/null +++ b/query.py @@ -0,0 +1,848 @@ +#!/usr/bin/env python3 + +import argparse +import csv +import json +import os +import re +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Set + +import requests + + +OSV_BATCH_URL = "https://api.osv.dev/v1/querybatch" +NVD_CVE_API = "https://services.nvd.nist.gov/rest/json/cves/2.0" + +UBUNTU_CVE_RE = re.compile(r"^UBUNTU-(CVE-\d{4}-\d{4,})$") +CVE_RE = re.compile(r"^(CVE-\d{4}-\d{4,})$") + + +# ---------------------------- +# Package parsing +# ---------------------------- + +def strip_debian_arch_suffix(pkg_name: str) -> str: + """ + Convert names like: + libyaml-0-2:arm64 + + into: + libyaml-0-2 + """ + if ":" not in pkg_name: + return pkg_name + + name, suffix = pkg_name.rsplit(":", 1) + + known_arches = { + "amd64", + "arm64", + "armhf", + "i386", + "ppc64el", + "s390x", + "riscv64", + "all", + } + + if suffix in known_arches: + return name + + return pkg_name + + +def parse_package_file(path: Path) -> List[Dict[str, str]]: + """ + Parses package files with rows like: + + package<TAB>version<TAB>arch + + Example: + + libyaml-0-2:arm64 0.2.2-1build2 arm64 + login 1:4.8.1-2ubuntu2.2 arm64 + """ + packages = [] + + with path.open("r", encoding="utf-8") as f: + for line_no, line in enumerate(f, start=1): + line = line.strip() + + if not line or line.startswith("#"): + continue + + parts = line.split() + + if len(parts) < 2: + print(f"[WARN] Skipping malformed line {line_no}: {line}") + continue + + raw_name = parts[0] + version = parts[1] + arch = parts[2] if len(parts) >= 3 else "" + + packages.append({ + "package_raw": raw_name, + "package": strip_debian_arch_suffix(raw_name), + "version": version, + "arch": arch, + }) + + return packages + + +def chunked(items: List[Any], size: int): + for i in range(0, len(items), size): + yield i, items[i:i + size] + + +# ---------------------------- +# OSV logic +# ---------------------------- + +def extract_fixed_versions(affected: List[Dict[str, Any]]) -> List[str]: + fixed = set() + + for affected_entry in affected: + for range_entry in affected_entry.get("ranges", []): + for event in range_entry.get("events", []): + if "fixed" in event: + fixed.add(event["fixed"]) + + return sorted(fixed) + + +def derive_cve_from_osv(osv_id: str, aliases: List[str]) -> str: + """ + Prefer actual alias CVE. + + If OSV ID looks like: + UBUNTU-CVE-2025-1149 + + derive: + CVE-2025-1149 + """ + for alias in aliases: + if CVE_RE.match(alias): + return alias + + match = UBUNTU_CVE_RE.match(osv_id) + if match: + return match.group(1) + + return "" + + +def ubuntu_cve_url(cve_id: str) -> str: + if cve_id.startswith("CVE-"): + return f"https://ubuntu.com/security/{cve_id}" + return "" + + +def query_osv( + packages: List[Dict[str, str]], + ecosystem: str, + batch_size: int, + sleep_seconds: float, +) -> List[Dict[str, Any]]: + rows = [] + + for start_idx, batch in chunked(packages, batch_size): + print( + f"[INFO] Querying OSV packages " + f"{start_idx + 1}-{start_idx + len(batch)} of {len(packages)}" + ) + + queries = [ + { + "package": { + "name": pkg["package"], + "ecosystem": ecosystem, + }, + "version": pkg["version"], + } + for pkg in batch + ] + + response = requests.post( + OSV_BATCH_URL, + json={"queries": queries}, + timeout=90, + ) + response.raise_for_status() + + results = response.json().get("results", []) + + if len(results) != len(batch): + print( + f"[WARN] OSV returned {len(results)} results " + f"for {len(batch)} queries" + ) + + for pkg, result in zip(batch, results): + for vuln in result.get("vulns", []): + osv_id = vuln.get("id", "") + aliases = vuln.get("aliases", []) + cves = [a for a in aliases if a.startswith("CVE-")] + + derived_cve = derive_cve_from_osv(osv_id, aliases) + + if not cves and derived_cve: + cves = [derived_cve] + + severity_entries = vuln.get("severity", []) + osv_severity = "; ".join( + f"{s.get('type', '')}:{s.get('score', '')}" + for s in severity_entries + ) + + fixed_versions = extract_fixed_versions(vuln.get("affected", [])) + + rows.append({ + "package_raw": pkg["package_raw"], + "package": pkg["package"], + "version": pkg["version"], + "arch": pkg["arch"], + "ecosystem": ecosystem, + + "osv_id": osv_id, + "cves": ",".join(cves), + "derived_cve": derived_cve, + "osv_summary": vuln.get("summary", ""), + "osv_details": vuln.get("details", ""), + "osv_severity": osv_severity, + "osv_published": vuln.get("published", ""), + "osv_modified": vuln.get("modified", ""), + "osv_fixed_versions": ",".join(fixed_versions), + "ubuntu_cve_url": ubuntu_cve_url(derived_cve), + "osv_references": ";".join( + ref.get("url", "") + for ref in vuln.get("references", []) + if ref.get("url") + ), + }) + + time.sleep(sleep_seconds) + + return rows + + +# ---------------------------- +# NVD parsing +# ---------------------------- + +def parse_cvss_vector(vector: str) -> Dict[str, str]: + """ + Parse vector strings like: + + CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H + + into: + + { + "AV": "N", + "AC": "L", + "PR": "N", + ... + } + """ + fields = {} + + if not vector: + return fields + + for part in vector.split("/"): + if ":" not in part: + continue + + key, value = part.split(":", 1) + + if key == "CVSS": + continue + + fields[key] = value + + return fields + + +def choose_best_cvss(metrics: Dict[str, Any]) -> Dict[str, Any]: + """ + Prefer newer CVSS versions. + + NVD may contain: + metrics.cvssMetricV40 + metrics.cvssMetricV31 + metrics.cvssMetricV30 + metrics.cvssMetricV2 + """ + priority = [ + ("cvssMetricV40", "4.0"), + ("cvssMetricV31", "3.1"), + ("cvssMetricV30", "3.0"), + ("cvssMetricV2", "2.0"), + ] + + for metric_key, version in priority: + entries = metrics.get(metric_key, []) + + if not entries: + continue + + selected = None + + for entry in entries: + if entry.get("source") == "nvd@nist.gov": + selected = entry + break + + if selected is None: + selected = entries[0] + + cvss_data = selected.get("cvssData", {}) + + return { + "nvd_cvss_version": version, + "nvd_vector": cvss_data.get("vectorString", ""), + "nvd_base_score": cvss_data.get("baseScore", ""), + "nvd_severity": ( + cvss_data.get("baseSeverity") + or selected.get("baseSeverity") + or "" + ), + "nvd_metric_source": selected.get("source", ""), + "nvd_metric_type": selected.get("type", ""), + } + + return empty_nvd_cvss_fields() + + +def empty_nvd_cvss_fields() -> Dict[str, Any]: + return { + "nvd_cvss_version": "", + "nvd_vector": "", + "nvd_base_score": "", + "nvd_severity": "", + "nvd_metric_source": "", + "nvd_metric_type": "", + } + + +def extract_english_description(cve_obj: Dict[str, Any]) -> str: + for desc in cve_obj.get("descriptions", []): + if desc.get("lang") == "en": + return desc.get("value", "") + + return "" + + +def parse_nvd_vulnerability(vuln: Dict[str, Any]) -> Dict[str, Any]: + cve_obj = vuln.get("cve", {}) + cve_id = cve_obj.get("id", "") + + cvss = choose_best_cvss(cve_obj.get("metrics", {})) + vector_fields = parse_cvss_vector(cvss.get("nvd_vector", "")) + + return { + "cve_id": cve_id, + "nvd_found": "yes", + "nvd_error": "", + "nvd_description": extract_english_description(cve_obj), + "nvd_published": cve_obj.get("published", ""), + "nvd_last_modified": cve_obj.get("lastModified", ""), + "nvd_vuln_status": cve_obj.get("vulnStatus", ""), + **cvss, + "nvd_AV": vector_fields.get("AV", ""), + "nvd_PR": vector_fields.get("PR", ""), + } + + +def empty_nvd_record(error: str = "") -> Dict[str, Any]: + return { + "nvd_found": "no", + "nvd_error": error, + "nvd_description": "", + "nvd_published": "", + "nvd_last_modified": "", + "nvd_vuln_status": "", + "nvd_cvss_version": "", + "nvd_vector": "", + "nvd_base_score": "", + "nvd_severity": "", + "nvd_metric_source": "", + "nvd_metric_type": "", + "nvd_AV": "", + "nvd_PR": "", + } + + +# ---------------------------- +# NVD API +# ---------------------------- + +def cve_year(cve_id: str) -> str: + """ + CVE-2025-1149 -> 2025 + """ + parts = cve_id.split("-") + if len(parts) >= 3: + return parts[1] + return "" + + +def nvd_headers(api_key: Optional[str]) -> Dict[str, str]: + if api_key: + return {"apiKey": api_key} + return {} + + +def fetch_nvd_cve_one_by_one( + cve_id: str, + api_key: Optional[str], + sleep_seconds: float, +) -> Dict[str, Any]: + response = requests.get( + NVD_CVE_API, + params={"cveId": cve_id}, + headers=nvd_headers(api_key), + timeout=90, + ) + + if response.status_code == 429: + raise RuntimeError( + "NVD rate limit hit. Use NVD_API_KEY or increase --nvd-sleep." + ) + + response.raise_for_status() + time.sleep(sleep_seconds) + + data = response.json() + vulns = data.get("vulnerabilities", []) + + if not vulns: + return empty_nvd_record("CVE not found in NVD") + + return parse_nvd_vulnerability(vulns[0]) + + +def fetch_nvd_year_page( + year: str, + start_index: int, + results_per_page: int, + api_key: Optional[str], + sleep_seconds: float, +) -> Dict[str, Any]: + params = { + "pubStartDate": f"{year}-01-01T00:00:00.000", + "pubEndDate": f"{year}-12-31T23:59:59.999", + "startIndex": start_index, + "resultsPerPage": results_per_page, + } + + response = requests.get( + NVD_CVE_API, + params=params, + headers=nvd_headers(api_key), + timeout=120, + ) + + if response.status_code == 429: + raise RuntimeError( + "NVD rate limit hit. Use NVD_API_KEY or increase --nvd-sleep." + ) + + response.raise_for_status() + time.sleep(sleep_seconds) + + return response.json() + + +def fetch_nvd_by_years( + needed_cves: List[str], + api_key: Optional[str], + sleep_seconds: float, + results_per_page: int, +) -> Dict[str, Dict[str, Any]]: + """ + Batch method. + + It fetches all NVD CVEs published in the years represented by the CVE IDs, + then keeps only the CVEs we need. + + Example: + needed CVEs: CVE-2025-1149, CVE-2025-1150, CVE-2022-27943 + + Fetch: + all NVD CVEs published in 2025 + all NVD CVEs published in 2022 + + This is much faster than one API call per CVE. + """ + needed: Set[str] = set(needed_cves) + years = sorted({cve_year(cve) for cve in needed if cve_year(cve)}) + + nvd_map: Dict[str, Dict[str, Any]] = {} + + print(f"[INFO] Fetching NVD in year batches: {', '.join(years)}") + + for year in years: + start_index = 0 + + while True: + print( + f"[INFO] NVD batch year={year}, " + f"startIndex={start_index}, resultsPerPage={results_per_page}" + ) + + data = fetch_nvd_year_page( + year=year, + start_index=start_index, + results_per_page=results_per_page, + api_key=api_key, + sleep_seconds=sleep_seconds, + ) + + total_results = data.get("totalResults", 0) + vulns = data.get("vulnerabilities", []) + + for vuln in vulns: + parsed = parse_nvd_vulnerability(vuln) + cve_id = parsed.get("cve_id", "") + + if cve_id in needed: + nvd_map[cve_id] = parsed + + start_index += results_per_page + + if start_index >= total_results: + break + + return nvd_map + + +# ---------------------------- +# Cache +# ---------------------------- + +def load_cache(path: Path) -> Dict[str, Dict[str, Any]]: + if path.exists(): + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + return {} + + +def save_cache(path: Path, cache: Dict[str, Dict[str, Any]]) -> None: + with path.open("w", encoding="utf-8") as f: + json.dump(cache, f, indent=2, ensure_ascii=False) + + +# ---------------------------- +# Enrichment +# ---------------------------- + +def enrich_with_nvd( + rows: List[Dict[str, Any]], + cache_path: Path, + nvd_sleep_seconds: float, + results_per_page: int, + use_batch: bool, + fallback_one_by_one: bool, +) -> List[Dict[str, Any]]: + api_key = "d753aa27-a3ae-4566-b52b-b3128208db21" + cache = load_cache(cache_path) + + unique_cves = sorted({ + row.get("derived_cve", "") + for row in rows + if row.get("derived_cve", "") + }) + + missing_cves = [ + cve for cve in unique_cves + if cve not in cache + ] + + print(f"[INFO] Unique CVEs to enrich with NVD: {len(unique_cves)}") + print(f"[INFO] Missing from local NVD cache: {len(missing_cves)}") + print(f"[INFO] NVD API key present: {'yes' if api_key else 'no'}") + print(f"[INFO] NVD mode: {'batch by year' if use_batch else 'one by one'}") + + if missing_cves: + if use_batch: + try: + batch_results = fetch_nvd_by_years( + needed_cves=missing_cves, + api_key=api_key, + sleep_seconds=nvd_sleep_seconds, + results_per_page=results_per_page, + ) + + cache.update(batch_results) + save_cache(cache_path, cache) + + except Exception as e: + print(f"[WARN] NVD batch fetch failed: {e}") + + still_missing = [ + cve for cve in missing_cves + if cve not in cache + ] + + if still_missing and fallback_one_by_one: + print( + f"[INFO] Falling back to one-by-one NVD lookup for " + f"{len(still_missing)} CVEs" + ) + + for idx, cve_id in enumerate(still_missing, start=1): + print( + f"[INFO] NVD fallback {idx}/{len(still_missing)}: {cve_id}" + ) + + try: + cache[cve_id] = fetch_nvd_cve_one_by_one( + cve_id=cve_id, + api_key=api_key, + sleep_seconds=nvd_sleep_seconds, + ) + except Exception as e: + cache[cve_id] = empty_nvd_record(str(e)) + + save_cache(cache_path, cache) + + elif still_missing: + for cve_id in still_missing: + cache[cve_id] = empty_nvd_record( + "Not found by batch method; fallback disabled" + ) + + save_cache(cache_path, cache) + + enriched_rows = [] + + for row in rows: + cve_id = row.get("derived_cve", "") + nvd_data = cache.get(cve_id, {}) if cve_id else {} + + new_row = dict(row) + + for key in [ + "nvd_found", + "nvd_error", + "nvd_severity", + "nvd_base_score", + "nvd_cvss_version", + "nvd_vector", + "nvd_AV", + "nvd_PR", + "nvd_metric_source", + "nvd_metric_type", + "nvd_vuln_status", + "nvd_published", + "nvd_last_modified", + "nvd_description", + ]: + new_row[key] = nvd_data.get(key, "") + + enriched_rows.append(new_row) + + return enriched_rows + + +# ---------------------------- +# Output +# ---------------------------- + +def write_csv(rows: List[Dict[str, Any]], path: Path) -> None: + fieldnames = [ + "package_raw", + "package", + "version", + "arch", + "ecosystem", + + "osv_id", + "cves", + "derived_cve", + "osv_summary", + "osv_severity", + "osv_published", + "osv_modified", + "osv_fixed_versions", + "ubuntu_cve_url", + + "nvd_found", + "nvd_error", + "nvd_severity", + "nvd_base_score", + "nvd_cvss_version", + "nvd_vector", + "nvd_AV", + "nvd_PR", + "nvd_metric_source", + "nvd_metric_type", + "nvd_vuln_status", + "nvd_published", + "nvd_last_modified", + "nvd_description", + + "osv_references", + "osv_details", + ] + + with path.open("w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=fieldnames, + extrasaction="ignore", + ) + writer.writeheader() + + for row in rows: + writer.writerow(row) + + +def write_json(rows: List[Dict[str, Any]], path: Path) -> None: + with path.open("w", encoding="utf-8") as f: + json.dump(rows, f, indent=2, ensure_ascii=False) + + +# ---------------------------- +# Main +# ---------------------------- + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Query OSV for vulnerable Ubuntu packages, then enrich CVEs " + "with NVD severity, CVSS vector, AV, and PR." + ) + ) + + parser.add_argument( + "input", + type=Path, + help="Input dpkg package list TSV", + ) + + parser.add_argument( + "--ecosystem", + default="Ubuntu:22.04:LTS", + help="OSV ecosystem, e.g. Ubuntu:22.04:LTS", + ) + + parser.add_argument( + "--csv", + type=Path, + default=Path("osv-nvd-results.csv"), + help="Output CSV path", + ) + + parser.add_argument( + "--json", + type=Path, + default=Path("osv-nvd-results.json"), + help="Output JSON path", + ) + + parser.add_argument( + "--nvd-cache", + type=Path, + default=Path("nvd-cache.json"), + help="Local NVD cache file path", + ) + + parser.add_argument( + "--osv-batch-size", + type=int, + default=1000, + help="Number of packages per OSV batch query", + ) + + parser.add_argument( + "--osv-sleep", + type=float, + default=0.2, + help="Seconds to sleep between OSV batch requests", + ) + + parser.add_argument( + "--nvd-sleep", + type=float, + default=None, + help=( + "Seconds to sleep between NVD requests. " + "Default: 6.0 without NVD_API_KEY, 0.7 with NVD_API_KEY." + ), + ) + + parser.add_argument( + "--nvd-results-per-page", + type=int, + default=2000, + help="NVD resultsPerPage for year-based batch fetching", + ) + + parser.add_argument( + "--skip-nvd", + action="store_true", + help="Only query OSV; do not enrich with NVD", + ) + + parser.add_argument( + "--no-nvd-batch", + action="store_true", + help="Disable year-based NVD batch fetching", + ) + + parser.add_argument( + "--no-nvd-fallback", + action="store_true", + help="Disable one-by-one fallback for CVEs missed by NVD batch fetch", + ) + + args = parser.parse_args() + + if args.nvd_sleep is None: + nvd_sleep_seconds = 0.7 if os.environ.get("NVD_API_KEY") else 6.0 + else: + nvd_sleep_seconds = args.nvd_sleep + + packages = parse_package_file(args.input) + + print(f"[INFO] Loaded packages: {len(packages)}") + print(f"[INFO] OSV ecosystem: {args.ecosystem}") + + rows = query_osv( + packages=packages, + ecosystem=args.ecosystem, + batch_size=args.osv_batch_size, + sleep_seconds=args.osv_sleep, + ) + + print(f"[INFO] OSV vulnerability rows: {len(rows)}") + + if not args.skip_nvd: + rows = enrich_with_nvd( + rows=rows, + cache_path=args.nvd_cache, + nvd_sleep_seconds=nvd_sleep_seconds, + results_per_page=args.nvd_results_per_page, + use_batch=not args.no_nvd_batch, + fallback_one_by_one=not args.no_nvd_fallback, + ) + + write_csv(rows, args.csv) + write_json(rows, args.json) + + print(f"[INFO] Wrote CSV: {args.csv}") + print(f"[INFO] Wrote JSON: {args.json}") + + if not args.skip_nvd: + print(f"[INFO] Wrote NVD cache: {args.nvd_cache}") + + +if __name__ == "__main__": + main() |