#!/usr/bin/env python3
"""Fold curated exploit archives and Vulhub environments 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 datetime import date, datetime, timezone
from pathlib import Path
from urllib.parse import urlsplit

import releases
from update_cves import ensure_cve_entries

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 publication_date(value: str, *, today: date | None = None) -> str | None:
    """ExploitDB supplies a calendar publication date, never an observation date."""
    if not re.fullmatch(r"\d{4}-\d{2}-\d{2}", value):
        return None
    try:
        published = date.fromisoformat(value)
    except ValueError:
        return None
    return value if published <= (today or datetime.now(timezone.utc).date()) else None


def exploitdb(publication_dates: dict[str, str] | None = None) -> 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 ""
        try:
            entry = int(row["id"])
        except (KeyError, TypeError, ValueError):
            continue
        if publication_dates is not None and (published := publication_date(str(row.get("date_published") or "").strip())):
            publication_dates[f"{EDB_ENTRY}{entry}"] = published
        if "CVE-" not in codes.upper():
            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:
    import link_health
    urls = [url for url in urls if not link_health.unavailable(url)]
    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 artifact_publication_date(url: str, publication_dates: dict[str, str]) -> str | None:
    parsed = urlsplit(url)
    if parsed.hostname not in {"exploit-db.com", "www.exploit-db.com"}:
        return None
    match = re.fullmatch(r"/exploits/(\d+)/?", parsed.path)
    return publication_dates.get(EDB_ENTRY + match[1]) if match else None


def sync_release_dates(publication_dates: dict[str, str], *, import_mode: bool = False) -> dict:
    pairs = releases.published_pairs()
    evidence = {releases.key(cve, url): {"source": published}
                for cve, url in pairs if (published := artifact_publication_date(url, publication_dates))}
    before = releases.load_ledger()
    ledger = releases.reconcile(before, pairs, evidence=evidence, import_mode=import_mode)
    releases.save_ledger(ledger)
    dated = [ledger[identity] for identity in evidence if ledger[identity].get("released")]
    return {"published_pairs": len(pairs), "edb_pairs_with_valid_source_date": len(evidence),
            "dated_edb_pairs": len(dated), "dated_edb_cves": len({row["cve"] for row in dated}),
            "new_release_rows": len(ledger) - len(before), "import_mode": import_mode}


def main() -> int:
    parser = argparse.ArgumentParser(description="Sync ExploitDB, Metasploit and Vulhub into the index")
    parser.add_argument("--dry-run", action="store_true", help="report without writing")
    parser.add_argument("--backfill-release-dates", action="store_true", help="Date existing ExploitDB links only; preserve historical import provenance")
    parser.add_argument("--cvelist-dir", type=Path, help="local CVE List V5 root for a large backfill")
    args = parser.parse_args()

    publication_dates: dict[str, str] = {}
    edb = exploitdb(publication_dates)
    if args.backfill_release_dates:
        if not args.dry_run:
            print(json.dumps(sync_release_dates(publication_dates, import_mode=True), sort_keys=True))
        return 0
    collected = [edb, metasploit(), vulhub()]
    created, unavailable = ensure_cve_entries(
        {cve for found in collected for cve in found},
        dry_run=args.dry_run,
        cvelist_dir=args.cvelist_dir,
    )
    if created or unavailable:
        verb = "would create" if args.dry_run else "created"
        print(f"base entries: {len(created):,} {verb}, {len(unavailable):,} unpublished or unavailable")

    for (header, _), found in zip(SOURCES, collected):
        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")
    if not args.dry_run:
        print(json.dumps(sync_release_dates(publication_dates), sort_keys=True))
    return 0


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