import argparse
import json
import os
import time
import zipfile
from io import BytesIO
from pathlib import Path

import urllib.request
import urllib.parse

BAZAAR_API = "https://mb-api.abuse.ch/api/v1/"
ZIP_PASSWORD = b"infected"

DATASET_DIR = Path(__file__).resolve().parent.parent / "datasets"
MALICIOUS_DIR = DATASET_DIR / "raw" / "malicious"

PE_TAGS = [
    "AgentTesla", "AsyncRAT", "RedLine", "Emotet", "NjRAT",
    "FormBook", "Remcos", "LokiBot", "GuLoader", "Raccoon",
    "Cobalt Strike", "Metasploit", "Sliver", "BruteRatel",
    "Ransomware", "Loader", "Dropper",
]

_API_KEY: str | None = None


def _post(payload: dict) -> dict:
    data = urllib.parse.urlencode(payload).encode()
    req = urllib.request.Request(BAZAAR_API, data=data, method="POST")
    req.add_header("User-Agent", "VibeCheck-Fetcher/1.0")
    if _API_KEY:
        req.add_header("API-KEY", _API_KEY)
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read())


def _download_sample(sha256: str, dest: Path) -> bool:
    payload = {"query": "get_file", "sha256_hash": sha256}
    data = urllib.parse.urlencode(payload).encode()
    req = urllib.request.Request(BAZAAR_API, data=data, method="POST")
    req.add_header("User-Agent", "VibeCheck-Fetcher/1.0")
    if _API_KEY:
        req.add_header("API-KEY", _API_KEY)
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            raw = resp.read()
        if raw[:2] != b"PK":
            return False
        with zipfile.ZipFile(BytesIO(raw)) as zf:
            zf.extractall(path=dest.parent, pwd=ZIP_PASSWORD)
        extracted = [f for f in dest.parent.iterdir() if f.suffix.lower() in {".exe", ".dll", ".sys", ".bin"} and f.name != dest.name]
        if extracted:
            extracted[0].rename(dest)
        return dest.exists()
    except Exception as e:
        print(f"  Download failed for {sha256[:16]}...: {e}")
        return False


def get_recent_samples(count: int) -> list[dict]:
    selector = min(count, 100)
    result = _post({"query": "get_recent", "selector": str(selector)})
    if result.get("query_status") != "ok":
        return []
    return result.get("data", [])


def get_tagged_samples(tag: str, limit: int = 100) -> list[dict]:
    result = _post({"query": "get_taginfo", "tag": tag, "limit": str(min(limit, 1000))})
    if result.get("query_status") not in ("ok", "tag_not_found"):
        return []
    return result.get("data") or []


def fetch(count: int, tags: list[str], pe_only: bool, resume: bool, rate_limit: float):
    MALICIOUS_DIR.mkdir(parents=True, exist_ok=True)

    existing = {f.stem for f in MALICIOUS_DIR.iterdir() if f.is_file()} if resume else set()
    print(f"  Resume mode: {len(existing)} samples already present")

    candidates: list[dict] = []

    if tags:
        for tag in tags:
            print(f"  Querying tag: {tag}")
            samples = get_tagged_samples(tag, limit=count)
            candidates.extend(samples)
            print(f"    -> {len(samples)} results")
            time.sleep(0.5)
    else:
        print(f"  Querying recent samples (up to {count})...")
        batches = (count + 99) // 100
        for _ in range(batches):
            samples = get_recent_samples(100)
            candidates.extend(samples)
            time.sleep(1.0)

    seen = set()
    deduped = []
    for s in candidates:
        h = s.get("sha256_hash", "")
        if h and h not in seen:
            seen.add(h)
            deduped.append(s)

    if pe_only:
        deduped = [s for s in deduped if s.get("file_type", "").lower() in {"exe", "dll", "sys"}]

    print(f"\n  {len(deduped)} unique PE candidates after filtering")

    downloaded = 0
    skipped = 0
    failed = 0

    for sample in deduped:
        if downloaded >= count:
            break

        sha256 = sample.get("sha256_hash", "")
        if not sha256:
            continue

        if sha256 in existing:
            skipped += 1
            continue

        file_type = sample.get("file_type", "bin").lower()
        ext = {"exe": ".exe", "dll": ".dll", "sys": ".sys"}.get(file_type, ".bin")
        dest = MALICIOUS_DIR / f"{sha256}{ext}"

        tags_str = ", ".join(sample.get("tags") or [])
        print(f"  [{downloaded + 1}/{count}] {sha256[:16]}... ({file_type}) [{tags_str}]", end=" ")

        ok = _download_sample(sha256, dest)
        if ok:
            print("OK")
            downloaded += 1
            existing.add(sha256)
        else:
            print("FAILED")
            failed += 1

        time.sleep(rate_limit)

    print(f"\n  Done: {downloaded} downloaded, {skipped} skipped (resume), {failed} failed")
    print(f"  Total malicious samples in {MALICIOUS_DIR}: {len(list(MALICIOUS_DIR.iterdir()))}")


def main():
    parser = argparse.ArgumentParser(description="Fetch malware samples from MalwareBazaar")
    parser.add_argument("--count", type=int, default=1000,
                        help="Number of samples to download (default: 1000)")
    parser.add_argument("--tags", nargs="+", default=None,
                        help="Filter by MalwareBazaar tags (e.g. --tags Emotet AgentTesla)")
    parser.add_argument("--pe-only", action="store_true", default=True,
                        help="Only download PE files (exe/dll/sys) (default: True)")
    parser.add_argument("--no-pe-only", dest="pe_only", action="store_false")
    parser.add_argument("--resume", action="store_true", default=True,
                        help="Skip already-downloaded samples (default: True)")
    parser.add_argument("--no-resume", dest="resume", action="store_false")
    parser.add_argument("--rate-limit", type=float, default=1.0,
                        help="Seconds between downloads (default: 1.0, be polite)")
    args = parser.parse_args()

    tags = args.tags or []
    print(f"MalwareBazaar Fetcher")
    print(f"  Target: {args.count} samples")
    print(f"  Tags:   {tags or 'recent (no filter)'}")
    print(f"  PE only: {args.pe_only}")
    print(f"  Output: {MALICIOUS_DIR}\n")

    fetch(
        count=args.count,
        tags=tags,
        pe_only=args.pe_only,
        resume=args.resume,
        rate_limit=args.rate_limit,
    )


if __name__ == "__main__":
    main()