#!/usr/bin/env python3
"""Fetch OSV ecosystem exports and extract the malicious-package records.

Hound's index holds malicious packages only, never vulnerability
advisories — they mean opposite things, and conflating them once had
Hound report tokio as malware. See crates/hound-defs/src/osv.rs.

This tool does the coarse filter so `build-pack` does not have to walk a
quarter of a million files it will discard. The Rust parser still applies
its own classifier, so a record slipping through here changes nothing.

    tools/ingest-osv.py <work-dir> [ecosystem ...]

Leaves <work-dir>/<ecosystem>/*.json ready for build-pack.
"""
import json
import os
import sys
import urllib.request
import zipfile

BASE = "https://osv-vulnerabilities.storage.googleapis.com"
DEFAULT = ["npm", "PyPI", "crates.io", "RubyGems", "Go", "Packagist"]


def is_malicious(record: dict) -> bool:
    """Mirror of osv::is_malicious_record in the Rust side."""
    if record.get("id", "").startswith("MAL-"):
        return True
    if "malicious-packages-origins" in record.get("database_specific", {}):
        return True
    summary = record.get("summary", "").lower()
    return summary.startswith("malicious code in") or summary.startswith("malicious package")


def fetch(eco: str, work: str) -> str:
    path = os.path.join(work, f"{eco}.zip")
    if os.path.exists(path) and os.path.getsize(path) > 0:
        print(f"{eco}: using cached {os.path.getsize(path) // 1048576} MB")
        return path
    url = f"{BASE}/{eco}/all.zip"
    print(f"{eco}: fetching {url}")
    urllib.request.urlretrieve(url, path)
    return path


def extract(eco: str, zip_path: str, work: str) -> int:
    out_dir = os.path.join(work, eco.replace(".", "-").lower())
    os.makedirs(out_dir, exist_ok=True)
    kept = skipped = 0
    with zipfile.ZipFile(zip_path) as z:
        for name in z.namelist():
            if not name.endswith(".json"):
                continue
            raw = z.read(name)
            try:
                record = json.loads(raw)
            except json.JSONDecodeError:
                skipped += 1
                continue
            if not is_malicious(record):
                skipped += 1
                continue
            with open(os.path.join(out_dir, os.path.basename(name)), "wb") as f:
                f.write(raw)
            kept += 1
    print(f"{eco}: {kept} malicious, {skipped} advisories skipped -> {out_dir}")
    return kept


def main() -> int:
    if len(sys.argv) < 2:
        print(__doc__)
        return 2
    work = sys.argv[1]
    ecosystems = sys.argv[2:] or DEFAULT
    os.makedirs(work, exist_ok=True)

    total = 0
    for eco in ecosystems:
        try:
            total += extract(eco, fetch(eco, work), work)
        except Exception as e:  # a missing export must not stop the rest
            print(f"{eco}: skipped ({e})")
    print(f"\ntotal malicious records: {total}")
    return 0


if __name__ == "__main__":
    sys.exit(main())