cve/scripts/sync_pocingithub.py

#!/usr/bin/env python3
"""Take candidate PoC repositories from PoC-in-GitHub and judge them here.

GitHub discovery only sees repositories pushed in the last few days, so a PoC
published years ago for an old CVE is unreachable no matter how long the job
runs. nomi-sec/PoC-in-GitHub has been collecting them daily since 1999 and
publishes one JSON file per CVE.

It collects without filtering, exactly as the workflow this index replaced
did, so nothing here is trusted on arrival: every candidate goes through
qualifying_repo_cves, the same gate a repository found by search must pass.
"""

from __future__ import annotations

import argparse
import io
import json
import os
import re
import sys
import tarfile
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from update_cves import (
    GITHUB_GRAPHQL_URL,
    CVES,
    GitHubClient,
    github_repo_from_url,
    http_json,
    is_blacklisted_repo,
    load_blacklist,
    qualifying_repo_cves,
    replace_section,
    section_links,
)

TARBALL = "https://codeload.github.com/nomi-sec/PoC-in-GitHub/tar.gz/refs/heads/master"
SECTION = "#### Github"
EMPTY = "No PoCs found on GitHub currently."
CVE_NAME = re.compile(r"^CVE-\d{4}-\d{4,}$", re.IGNORECASE)
BATCH = 10
USER_AGENT = "0xMarcio-cve-poc-in-github"

FIELDS = """
  nameWithOwner url description isFork isArchived
  repositoryTopics(first: 20) { nodes { topic { name } } }
  readmeMd: object(expression: "HEAD:README.md") { ... on Blob { text } }
  readmeUpper: object(expression: "HEAD:README.MD") { ... on Blob { text } }
  readmeRst: object(expression: "HEAD:README.rst") { ... on Blob { text } }
  readmeBare: object(expression: "HEAD:README") { ... on Blob { text } }
  root: object(expression: "HEAD:") { ... on Tree { entries { name type } } }
"""


def candidates() -> dict[str, set[str]]:
    """CVE to the repository URLs the collection lists for it."""
    request = urllib.request.Request(TARBALL, headers={"User-Agent": USER_AGENT})
    with urllib.request.urlopen(request, timeout=600) as response:
        archive = response.read()
    found: dict[str, set[str]] = {}
    with tarfile.open(fileobj=io.BytesIO(archive), mode="r:gz") as tar:
        for member in tar:
            if not member.isfile() or not member.name.endswith(".json"):
                continue
            cve = member.name.rsplit("/", 1)[-1][:-5].upper()
            if not CVE_NAME.match(cve):
                continue
            handle = tar.extractfile(member)
            if handle is None:
                continue
            try:
                listed = json.loads(handle.read())
            except (ValueError, UnicodeDecodeError):
                continue
            urls = {
                str(item.get("html_url") or "").rstrip("/")
                for item in listed
                if item.get("html_url")
            }
            if urls:
                found[cve] = urls
    return found


def already_linked() -> dict[str, tuple[Path, set[str]]]:
    """What each entry in this index already lists under GitHub."""
    held: dict[str, tuple[Path, set[str]]] = {}
    for path in sorted(CVES.glob("[12][0-9][0-9][0-9]/CVE-*.md")):
        text = path.read_text(encoding="utf-8", errors="replace")
        held[path.stem.upper()] = (path, {url.rstrip("/") for url in section_links(text, SECTION)})
    return held


def describe(client: GitHubClient, names: list[str]) -> dict[str, dict]:
    """Metadata, README and root listing for a batch, as the classifier expects."""
    aliases = " ".join(
        f"r{index}: repository(owner: {json.dumps(name.split('/', 1)[0])}, "
        f"name: {json.dumps(name.split('/', 1)[1])}) {{ {FIELDS} }}"
        for index, name in enumerate(names)
    )
    payload = http_json(
        GITHUB_GRAPHQL_URL,
        headers=client.headers,
        data={"query": "query { " + aliases + " rateLimit { remaining } }"},
    )
    data = payload.get("data") or {}
    described: dict[str, dict] = {}
    for index, name in enumerate(names):
        repo = data.get(f"r{index}")
        if repo:
            described[name] = repo
    return described


def main() -> int:
    parser = argparse.ArgumentParser(description="Judge PoC-in-GitHub candidates against this index")
    parser.add_argument("--limit", type=int, default=0, help="Repositories to judge this run")
    parser.add_argument("--dry-run", action="store_true", help="Report without writing files")
    args = parser.parse_args()

    listed = candidates()
    held = already_linked()
    print(f"{len(listed):,} CVEs listed upstream | {len(held):,} entries in this index")

    # One repository can be offered for several CVEs; judge it once per pair.
    blacklist = load_blacklist()
    pending: dict[str, set[str]] = {}
    for cve, urls in listed.items():
        entry = held.get(cve)
        if entry is None:
            continue
        for url in urls - entry[1]:
            full_name = github_repo_from_url(url)
            if full_name and "/" in full_name and not is_blacklisted_repo(full_name, blacklist):
                pending.setdefault(full_name, set()).add(cve)

    names = sorted(pending)
    if args.limit:
        names = names[: args.limit]
    print(f"{sum(len(pending[n]) for n in names):,} candidate links across {len(names):,} repositories")

    token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") or ""
    client = GitHubClient(token)
    accepted: dict[str, set[str]] = {}
    judged = kept = 0

    batches = [names[i : i + BATCH] for i in range(0, len(names), BATCH)]
    with ThreadPoolExecutor(max_workers=4) as executor:
        futures = {executor.submit(describe, client, batch): batch for batch in batches}
        for future in as_completed(futures):
            batch = futures[future]
            try:
                described = future.result()
            except Exception as problem:
                print(f"Skipped a batch starting at {batch[0]}: {problem}", file=sys.stderr)
                described = {}
            for full_name in batch:
                judged += 1
                repo = described.get(full_name)
                if repo is None:
                    continue
                for cve in pending[full_name]:
                    year = int(cve.split("-")[1])
                    if cve in qualifying_repo_cves(repo, year, blacklist):
                        accepted.setdefault(cve, set()).add(str(repo.get("url") or "").rstrip("/"))
                        kept += 1
            if judged % 2000 < BATCH:
                print(f"  judged {judged:,} of {len(names):,}")

    written = 0
    for cve, urls in sorted(accepted.items()):
        path, existing = held[cve]
        merged = section_links(path.read_text(encoding="utf-8", errors="replace"), SECTION)
        merged = merged + sorted(url for url in urls if url.rstrip("/") not in existing)
        with path.open("r", encoding="utf-8", newline="") as handle:
            text = handle.read()
        updated, changed = replace_section(text, SECTION, merged, EMPTY)
        if changed and not args.dry_run:
            with path.open("w", encoding="utf-8", newline="") as handle:
                handle.write(updated)
        written += bool(changed)

    print(f"accepted {kept:,} links for {len(accepted):,} CVEs | {written:,} entries updated")
    print(f"rejected {sum(len(pending[n]) for n in names) - kept:,} candidates that did not pass the gate")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())