#!/usr/bin/env python3
"""Fold the curated exploit archives into the index.

ExploitDB and Metasploit both publish a single machine readable file mapping
their entries to CVE ids. Neither is a GitHub repository, so their links are
kept in their own sections rather than being counted and starred as one.
"""

from __future__ import annotations

import argparse
import csv
import gzip
import io
import json
import re
import urllib.request
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
CVES = ROOT / "cves"
INDEX = ROOT / "index"
EDB_CSV = "https://gitlab.com/exploit-database/exploitdb/-/raw/main/files_exploits.csv"
MSF_JSON = ("https://raw.githubusercontent.com/rapid7/metasploit-framework/"
            "master/db/modules_metadata_base.json")
EDB_ENTRY = "https://www.exploit-db.com/exploits/"
MSF_BLOB = "https://github.com/rapid7/metasploit-framework/blob/master"
CVE = re.compile(r"CVE-\d{4}-\d{4,7}", re.I)


def plausible(cve: str) -> bool:
    """Reject impossible years. Metasploit ships CVE-1012-5358 for CVE-2012-5358."""
    return 1999 <= int(cve.split("-")[1]) <= 2100
USER_AGENT = "0xMarcio-cve-exploits"
# One CVE can carry dozens of ExploitDB rows for the same flaw on different
# platforms. Enough to be useful, not enough to bury the repositories below.
PER_CVE = 8

SOURCES = (
    ("#### ExploitDB", "No ExploitDB entry."),
    ("#### Metasploit", "No Metasploit module."),
    ("#### Vulhub", "No Vulhub environment."),
)

VULHUB_TARBALL = "https://codeload.github.com/vulhub/vulhub/tar.gz/refs/heads/master"
VULHUB_TREE = "https://github.com/vulhub/vulhub/tree/master/"


def fetch(url: str) -> bytes:
    request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    with urllib.request.urlopen(request, timeout=180) as response:
        raw = response.read()
    return gzip.decompress(raw) if url.endswith(".gz") else raw


def exploitdb() -> dict[str, list[str]]:
    """CVE to ExploitDB entry links, verified entries first, newest first."""
    text = fetch(EDB_CSV).decode("utf-8", "replace")
    rows: dict[str, list[tuple[int, int]]] = defaultdict(list)
    for row in csv.DictReader(io.StringIO(text)):
        codes = row.get("codes") or ""
        if "CVE-" not in codes.upper():
            continue
        try:
            entry = int(row["id"])
        except (KeyError, TypeError, ValueError):
            continue
        verified = 1 if str(row.get("verified") or "").strip() == "1" else 0
        for cve in {c.upper() for c in CVE.findall(codes)}:
            rows[cve].append((verified, entry))
    return {
        cve: [f"{EDB_ENTRY}{entry}" for _, entry in
              sorted(found, key=lambda item: (-item[0], -item[1]))[:PER_CVE]]
        for cve, found in rows.items()
    }


def metasploit() -> dict[str, list[str]]:
    """CVE to module source links, the better ranked module first.

    A module is written to be run, so rank is the useful order here: excellent
    beats manual, whatever order the database happens to list them in.
    """
    data = json.loads(fetch(MSF_JSON).decode("utf-8", "replace"))
    rows: dict[str, list[tuple[int, str]]] = defaultdict(list)
    for module in data.values():
        path = str(module.get("path") or "")
        if not path.endswith(".rb"):
            continue
        rank = int(module.get("rank") or 0)
        found: set[str] = set()
        for reference in module.get("references") or []:
            found |= {c.upper() for c in CVE.findall(str(reference))}
        found = {cve for cve in found if plausible(cve)}
        # An exploit module is written against the flaws it lists, chains
        # included. A scanner is not: "SSL/TLS Version Detection" cites BEAST,
        # RC4, POODLE and Logjam because it reports on them, and filing it as a
        # proof of concept under each is wrong. Keep auxiliary and post modules
        # only where they name a single CVE, which is the shape of an auxiliary
        # module that is itself the exploit.
        if str(module.get("type") or "") != "exploit" and len(found) > 1:
            continue
        for cve in found:
            rows[cve].append((rank, MSF_BLOB + path))
    return {
        cve: [url for _, url in
              sorted(found, key=lambda item: (-item[0], item[1]))[:PER_CVE]]
        for cve, found in rows.items()
    }


def vulhub() -> dict[str, list[str]]:
    """CVE to the directory holding a runnable vulnerable environment.

    Vulhub ships a docker-compose file and reproduction steps per flaw, which
    is a proof of concept anyone can stand up rather than read about. One
    directory per environment, the CVE named in its path.
    """
    import io
    import tarfile

    rows: dict[str, set[str]] = defaultdict(set)
    with tarfile.open(fileobj=io.BytesIO(fetch(VULHUB_TARBALL)), mode="r:gz") as tar:
        for member in tar:
            parts = Path(member.name).parts
            # <archive>/<product>/<environment>; anything deeper is its contents
            if len(parts) != 3 or not member.isdir():
                continue
            for cve in {c.upper() for c in CVE.findall(parts[-1])}:
                if plausible(cve):
                    rows[cve].add(VULHUB_TREE + "/".join(parts[1:]))
    return {cve: sorted(found)[:PER_CVE] for cve, found in rows.items()}


def section_bounds(text: str, header: str) -> tuple[int, int] | None:
    start = text.find("\n" + header)
    if start < 0:
        return None
    body = start + 1 + len(header)
    nxt = text.find("\n#", body)
    return body, len(text) if nxt < 0 else nxt


def apply_to_markdown(cve: str, header: str, urls: list[str], dry_run: bool) -> str:
    path = CVES / cve.split("-")[1] / f"{cve}.md"
    if not path.exists():
        return "absent"
    original = path.read_text(encoding="utf-8")
    bounds = section_bounds(original, header)
    block = "\n" + "\n".join(f"- {url}" for url in urls) + "\n"
    if bounds:
        start, end = bounds
        if original[start:end] == block:
            return "unchanged"
        updated = original[:start] + block + original[end:]
    else:
        updated = original.rstrip("\n") + f"\n\n{header}\n" + block.lstrip("\n")
    if not dry_run:
        path.write_text(updated, encoding="utf-8")
    return "added"


def drop_section(cve: str, header: str, dry_run: bool) -> bool:
    """Remove a section for a CVE the source no longer claims."""
    path = CVES / cve.split("-")[1] / f"{cve}.md"
    if not path.exists():
        return False
    original = path.read_text(encoding="utf-8")
    bounds = section_bounds(original, header)
    if not bounds:
        return False
    start, _ = bounds
    head = original.rfind("\n" + header, 0, start + 1)
    end = bounds[1]
    updated = (original[:head] + original[end:]).rstrip("\n") + "\n"
    if not dry_run:
        path.write_text(updated, encoding="utf-8")
    return True


def main() -> int:
    parser = argparse.ArgumentParser(description="Sync ExploitDB and Metasploit into the index")
    parser.add_argument("--dry-run", action="store_true", help="report without writing")
    args = parser.parse_args()

    for (header, _), collect in zip(SOURCES, (exploitdb, metasploit, vulhub)):
        found = collect()
        links = sum(len(urls) for urls in found.values())
        tally = {"added": 0, "unchanged": 0, "absent": 0}
        for cve, urls in sorted(found.items()):
            tally[apply_to_markdown(cve, header, urls, args.dry_run)] += 1
        # A CVE the source has stopped claiming is never visited above, so its
        # section would otherwise survive every future run.
        marker = "\n" + header
        stale = 0
        for path in CVES.glob("[12][0-9][0-9][0-9]/CVE-*.md"):
            if path.stem in found:
                continue
            if marker in path.read_text(encoding="utf-8", errors="replace"):
                stale += drop_section(path.stem, header, args.dry_run)
        name = header.replace("#", "").strip()
        print(f"{name}: {len(found):,} CVEs, {links:,} links | "
              f"{tally['added']:,} entries written, {tally['unchanged']:,} unchanged, "
              f"{tally['absent']:,} for CVEs not in the index, {stale:,} stale sections dropped")
    return 0


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