"""Pull DataDog malicious-software-packages-dataset → JSONL via parallel raw downloads.

No git clone. Uses GitHub tree API + raw.githubusercontent.com + ThreadPoolExecutor.

Usage:
    python scripts/pull_datadog.py --max-samples 400 --ecosystems npm pypi
"""
from __future__ import annotations

import argparse
import io
import json
import os
import re
import sys
import time
import zipfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import requests

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


REPO = "DataDog/malicious-software-packages-dataset"
RAW_BASE = f"https://raw.githubusercontent.com/{REPO}/main"
ZIP_PASSWORD = b"infected"

_INTERESTING_FILES = (
    "package.json", "setup.py", "setup.cfg", "pyproject.toml",
    "index.js", "preinstall.js", "postinstall.js", "install.js",
    "__init__.py",
)

_MALICIOUS_TOKENS = re.compile(
    r"(eval\(|exec\(|subprocess|os\.system|child_process|"
    r"base64\.b64decode|atob\(|Buffer\.from\([\"'][A-Za-z0-9+/=]{20,}|"
    r"urllib\.request\.url(retrieve|open)|requests\.(get|post|put)|fetch\(|"
    r"\.exec\(|new Function\(|"
    r"/etc/passwd|\.ssh/|\.aws/|\.npmrc|\.env|"
    r"hostname|whoami|process\.env|os\.environ|"
    r"verify_signature|verify=False|InsecureRequestWarning|"
    r"importlib\.import_module|require\([a-zA-Z_])",
    re.IGNORECASE,
)


def _gh_headers() -> dict:
    h = {"Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28"}
    tok = os.environ.get("GITHUB_TOKEN")
    if tok:
        h["Authorization"] = f"Bearer {tok}"
    return h


def fetch_tree() -> list[str]:
    """Fetch full repo tree (recursive). Returns list of paths to .zip files."""
    paths: list[str] = []
    url = f"https://api.github.com/repos/{REPO}/git/trees/main?recursive=1"
    r = requests.get(url, headers=_gh_headers(), timeout=30)
    r.raise_for_status()
    data = r.json()
    paths.extend(t["path"] for t in data.get("tree", []) if t["path"].endswith(".zip"))
    if data.get("truncated"):
        for category in ("malicious_intent", "compromised_lib"):
            sub_url = f"https://api.github.com/repos/{REPO}/contents/samples/pypi/{category}"
            sr = requests.get(sub_url, headers=_gh_headers(), timeout=30)
            if sr.status_code != 200:
                continue
            pkgs = sr.json() if isinstance(sr.json(), list) else []
            for pkg in pkgs[:80]:
                if pkg.get("type") != "dir":
                    continue
                vr = requests.get(pkg["url"], headers=_gh_headers(), timeout=30)
                if vr.status_code != 200:
                    continue
                vers = vr.json() if isinstance(vr.json(), list) else []
                for ver in vers[:2]:
                    if ver.get("type") != "dir":
                        continue
                    zr = requests.get(ver["url"], headers=_gh_headers(), timeout=30)
                    if zr.status_code != 200:
                        continue
                    items = zr.json() if isinstance(zr.json(), list) else []
                    for z in items:
                        if z.get("name", "").endswith(".zip"):
                            paths.append(z["path"])
    return paths


def stratify(paths: list[str], ecosystems: list[str], max_samples: int) -> list[str]:
    buckets: dict[tuple[str, str], list[str]] = {}
    for p in paths:
        parts = p.split("/")
        if len(parts) < 4 or parts[0] != "samples":
            continue
        eco = parts[1]
        cat = parts[2]
        if eco not in ecosystems:
            continue
        buckets.setdefault((eco, cat), []).append(p)
    n_buckets = max(1, len(buckets))
    per_bucket = max(1, max_samples // n_buckets)
    chosen: list[str] = []
    for (eco, cat), items in sorted(buckets.items()):
        chosen.extend(items[:per_bucket])
    return chosen[:max_samples]


def fetch_one(path: str, attempt_max: int = 3) -> tuple[str, bytes | None]:
    url = f"{RAW_BASE}/{path}"
    for attempt in range(attempt_max):
        try:
            r = requests.get(url, timeout=20)
            if r.status_code == 200:
                return path, r.content
            time.sleep(0.5 * (attempt + 1))
        except requests.RequestException:
            time.sleep(0.5 * (attempt + 1))
    return path, None


def extract_sample(path: str, blob: bytes) -> dict | None:
    parts = path.split("/")
    if len(parts) < 6:
        return None
    eco = parts[1]
    category = parts[2]
    package = parts[3]
    version = parts[4]
    fname = parts[-1]
    try:
        with zipfile.ZipFile(io.BytesIO(blob)) as zf:
            interesting: list[tuple[str, str]] = []
            for name in zf.namelist():
                base = os.path.basename(name).lower()
                if base in _INTERESTING_FILES or any(name.endswith(ext) for ext in (".js", ".py")):
                    try:
                        with zf.open(name, pwd=ZIP_PASSWORD) as fh:
                            raw = fh.read()
                            try:
                                text = raw.decode("utf-8", errors="ignore")
                            except Exception:
                                continue
                        interesting.append((name, text))
                    except (RuntimeError, zipfile.BadZipFile):
                        continue
            if not interesting:
                return None
            interesting.sort(key=lambda t: (-len(t[1]), t[0]))
            top_name, top_text = interesting[0]
            blob_for_scan = "\n".join(t[1] for t in interesting[:5])[:8000]
            n_signals = len(_MALICIOUS_TOKENS.findall(blob_for_scan))
            return {
                "ecosystem": eco,
                "category": category,
                "package": package,
                "version": version,
                "filename": fname,
                "primary_file": top_name,
                "diff_preview": top_text[:1500],
                "files": [n for n, _ in interesting[:6]],
                "malicious_signal_count": n_signals,
                "label": "malicious",
                "source": "datadog",
            }
    except (zipfile.BadZipFile, RuntimeError):
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--max-samples", type=int, default=400)
    ap.add_argument("--ecosystems", nargs="*", default=["npm", "pypi"])
    ap.add_argument("--workers", type=int, default=32)
    ap.add_argument("--out", default="data/datadog_extracted.jsonl")
    args = ap.parse_args()

    print(f"[tree] fetching tree from github...", flush=True)
    t0 = time.time()
    paths = fetch_tree()
    print(f"[tree] {len(paths)} zip paths discovered in {time.time()-t0:.1f}s", flush=True)

    chosen = stratify(paths, args.ecosystems, args.max_samples)
    print(f"[stratify] picked {len(chosen)} samples (target={args.max_samples})", flush=True)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    n_ok = 0
    n_fail = 0
    n_skipped = 0

    t1 = time.time()
    with open(out, "w", encoding="utf-8") as f, ThreadPoolExecutor(max_workers=args.workers) as ex:
        futures = [ex.submit(fetch_one, p) for p in chosen]
        for i, fut in enumerate(as_completed(futures), 1):
            path, blob = fut.result()
            if blob is None:
                n_fail += 1
            else:
                rec = extract_sample(path, blob)
                if rec is None:
                    n_skipped += 1
                else:
                    f.write(json.dumps(rec, ensure_ascii=False) + "\n")
                    n_ok += 1
            if i % 50 == 0:
                print(f"  progress {i}/{len(chosen)} (ok={n_ok}, fail={n_fail}, skip={n_skipped}) "
                      f"in {time.time()-t1:.1f}s", flush=True)

    elapsed = time.time() - t1
    print(f"DONE: {n_ok} extracted, {n_fail} fetch-fail, {n_skipped} no-interesting-files "
          f"in {elapsed:.1f}s -> {out}", flush=True)


if __name__ == "__main__":
    main()