hack/tools/scripts/llm-key-hunter/vault_import.py

#!/usr/bin/env python3
"""Import hunt results into the usable_keys vault.

Reads every results/<hunt>/*.txt bucket file for USABLE verdicts, merges
into usable_keys/<provider>/keys.json (dedup by key), and rebuilds the
per-provider keys.env + key_NN.txt files plus the root all_keys.json and
README.md index.

Only NEW keys are added — existing entries are left untouched (their
detail/models stay as first verified).

Usage:
    python3 vault_import.py [--dry-run] [--source results/ai_tools]

Bucket file shape (tab or pipe separated, first field = verdict):
    USABLE|key|url|detail
"""

import argparse
import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent
RESULTS = HERE / "results"
VAULT = HERE / "usable_keys"

# verdict label -> bucket filename inside a results/<hunt>/ dir
BUCKET_NAMES = {
    "USABLE": ("usable.txt", "usable_keys.txt", "live_keys_raw.txt",
               "all_non_401.txt"),
}

# provider name normalization for entries whose hunt dir doesn't match
# the vault dir name (e.g. ai_tools/codeium -> vault has no codeium dir yet)
PROVIDER_ALIASES = {
    "codeium": "codeium",
    "copilot": "copilot",
    "claude": "claude",
    "cline": "cline",
    "continue": "continue",
    "generic-env": "generic-env",
    "aws-sso": "aws-sso",
    "deepseek": "deepseek",
    "kimi": "moonshot",
    "moonshot": "moonshot",
    "exa": "exa",
    "codingplan": "codingplan",
    "freemodel": "freemodel",
    "xunfei": "xunfei",
    "kiro": "kiro",
    "minimax": "minimax",
    "siliconflow": "siliconflow",
    "volcanoark": "volcanoark",
    "zhipu": "zhipu",
    "mimo": "mimo",
    "dashscope": "dashscope",
    "longcat": "longcat",
}

MODELS_HINT = {
    "exa": ["exa-search"],
    "deepseek": ["deepseek-chat", "deepseek-reasoner"],
}


def parse_bucket_line(line):
    """Return (verdict, key, source, detail) from a bucket line."""
    for sep in ("|", "\t"):
        parts = line.split(sep)
        if len(parts) >= 2:
            return parts[0].strip(), parts[1].strip(), \
                (parts[2].strip() if len(parts) > 2 else ""), \
                (parts[3].strip() if len(parts) > 3 else "")
    return None, None, None, None


def load_vault(provider):
    """Load usable_keys/<provider>/keys.json -> list of dicts."""
    p = VAULT / provider / "keys.json"
    if not p.exists():
        return []
    try:
        data = json.loads(p.read_text())
        return data if isinstance(data, list) else [data]
    except Exception:
        return []


def guess_provider(hunt_dir, line_key):
    """Map a hunt dir + key to a vault provider name."""
    if hunt_dir in PROVIDER_ALIASES:
        return PROVIDER_ALIASES[hunt_dir]
    if hunt_dir.startswith("ai_tools"):
        return hunt_dir.split("/")[-1]
    return hunt_dir


def rebuild_provider(provider, entries):
    """Write keys.json, keys.env, key_NN.txt for one provider."""
    d = VAULT / provider
    d.mkdir(parents=True, exist_ok=True)
    (d / "keys.json").write_text(
        json.dumps(entries, indent=2, ensure_ascii=False) + "\n")

    # keys.env: KEY_<N>=<key>
    env_lines = []
    for i, e in enumerate(entries):
        env_lines.append(f"{provider.upper()}_KEY_{i + 1}={e['key']}")
    (d / "keys.env").write_text("\n".join(env_lines) + "\n")

    # per-key txt: key + source + detail
    for i, e in enumerate(entries):
        (d / f"key_{i + 1:02d}.txt").write_text(
            f"key: {e['key']}\nsource: {e.get('source', '')}\n"
            f"detail: {e.get('detail', '')}\n")


def rebuild_root():
    """Rebuild usable_keys/all_keys.json + README.md from all providers."""
    all_keys, counts = [], {}
    for prov_dir in sorted(p for p in VAULT.iterdir() if p.is_dir()):
        f = prov_dir / "keys.json"
        if not f.exists():
            continue
        try:
            arr = json.loads(f.read_text())
        except Exception:
            continue
        arr = arr if isinstance(arr, list) else [arr]
        for e in arr:
            e = dict(e)
            e.setdefault("provider", prov_dir.name)
            all_keys.append(e)
        counts[prov_dir.name] = len(arr)

    (VAULT / "all_keys.json").write_text(
        json.dumps(all_keys, indent=2, ensure_ascii=False) + "\n")

    total = sum(counts.values())
    lines = ["# Usable LLM Keys Vault", "", "## Counts", ""]
    lines += [f"- **{k}**: {v}" for k, v in
              sorted(counts.items(), key=lambda x: -x[1])]
    lines.append(f"\n**Total: {total} keys across {len(counts)} providers**")
    lines.append(f"\n_Updated: {datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%MZ')}_")
    (VAULT / "README.md").write_text("\n".join(lines) + "\n")
    return counts, total


def collect_usable(sources):
    """Scan results dirs for USABLE bucket lines. Returns {provider: {key: entry}}."""
    found = {}
    for src in sources:
        src = Path(src)
        if not src.is_dir():
            continue
        for bucket_name in BUCKET_NAMES["USABLE"]:
            bp = src / bucket_name
            if not bp.exists():
                continue
            for line in bp.read_text(errors="replace").splitlines():
                line = line.strip()
                if not line:
                    continue
                verdict, key, source, detail = parse_bucket_line(line)
                if verdict != "USABLE":
                    # all_non_401.txt has no verdict prefix: USABLE|deepseek|key|...
                    if "|" in line and line.split("|")[0] in (
                            "USABLE", "NO_BALANCE", "NO_ACCESS"):
                        continue
                    if line.startswith("LIVE|"):
                        _, prov, key, source, detail = line.split("|", 4)
                        provider = guess_provider(src.name, key)
                        found.setdefault(provider, {})[key] = {
                            "key": key, "source": source,
                            "detail": detail[:300], "provider": provider}
                    continue
                provider = guess_provider(src.name, key)
                found.setdefault(provider, {})[key] = {
                    "key": key, "source": source, "detail": (detail or "")[:300],
                    "provider": provider}
    return found


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--source", action="append", default=None,
                    help="results dir to scan (default: all results/*)")
    args = ap.parse_args()

    if args.source:
        sources = [RESULTS / s for s in args.source]
    else:
        sources = sorted(p for p in RESULTS.iterdir() if p.is_dir()
                         and p.name != "_cache")

    found = collect_usable(sources)
    if not found:
        print("no USABLE keys found in results dirs.")
        return

    added_total = 0
    report_lines = []
    for provider in sorted(found):
        existing = {e["key"] for e in load_vault(provider)}
        new = {k: v for k, v in found[provider].items() if k not in existing}
        if not new:
            continue
        entries = load_vault(provider) + [
            {**v, "base_url": v.get("base_url", ""),
             "auth": v.get("auth", "Bearer"),
             "models": v.get("models", MODELS_HINT.get(provider, []))}
            for v in new.values()]
        if not args.dry_run:
            rebuild_provider(provider, entries)
        added_total += len(new)
        report_lines.append(f"  {provider:14s}: +{len(new)} "
                            f"(vault now {len(entries)})")

    if not args.dry_run and added_total:
        counts, total = rebuild_root()
        report_lines.append(f"  {'TOTAL':14s}: +{added_total} "
                            f"(vault {total} keys, {len(counts)} providers)")

    print("\n".join(report_lines) if report_lines
          else "no new keys to add.")
    print(f"\n{'[DRY RUN] ' if args.dry_run else ''}added {added_total} keys.")


if __name__ == "__main__":
    main()