import argparse
import os
import random
import struct
import sys
from pathlib import Path

import httpx


def mutate_binary(original: bytes, mutation_id: int) -> bytes:
    data = bytearray(original)
    rng = random.Random(mutation_id)

    safe_start = min(512, len(data) - 1)
    for _ in range(rng.randint(5, 20)):
        pos = rng.randint(safe_start, len(data) - 1)
        data[pos] = rng.randint(0, 255)

    junk_size = rng.randint(16, 256)
    junk = bytes(rng.randint(0, 255) for _ in range(junk_size))
    data.extend(junk)

    if len(data) >= 0x3C + 4:
        pe_offset_bytes = data[0x3C:0x3C + 4]
        pe_offset = struct.unpack_from("<I", pe_offset_bytes)[0]
        checksum_offset = pe_offset + 88
        if checksum_offset + 4 <= len(data):
            struct.pack_into("<I", data, checksum_offset, rng.randint(0, 0xFFFFFFFF))

        timestamp_offset = pe_offset + 8
        if timestamp_offset + 4 <= len(data):
            struct.pack_into("<I", data, timestamp_offset, rng.randint(0, 0xFFFFFFFF))

    return bytes(data)


def run_validation(binary_path: str, api_url: str, mutation_count: int = 10):
    filepath = Path(binary_path)
    if not filepath.exists():
        print(f"File not found: {filepath}")
        sys.exit(1)

    original = filepath.read_bytes()
    print(f"Original binary: {filepath.name} ({len(original)} bytes)")
    print(f"API: {api_url}")
    print(f"Mutations: {mutation_count}")
    print()

    client = httpx.Client(timeout=30.0)
    results = []

    print("Submitting original...")
    resp = client.post(
        f"{api_url}/scan",
        files={"file": (filepath.name, original, "application/octet-stream")},
    )
    resp.raise_for_status()
    data = resp.json()
    results.append({
        "name": "original",
        "is_new": data["is_new_vibe"],
        "distance": data.get("tlsh_distance", 0),
        "verdict": data["final_verdict"],
    })

    for i in range(1, mutation_count + 1):
        mutated = mutate_binary(original, i)
        label = f"mut_{i:02d}"
        print(f"Submitting {label}...")
        resp = client.post(
            f"{api_url}/scan",
            files={"file": (f"{label}.exe", mutated, "application/octet-stream")},
        )
        resp.raise_for_status()
        data = resp.json()
        results.append({
            "name": label,
            "is_new": data["is_new_vibe"],
            "distance": data.get("tlsh_distance", 0) or 0,
            "verdict": data["final_verdict"],
        })

    client.close()

    print()
    print(f"{'Sample':<12} {'New Vibe?':<12} {'Distance':<12} {'Verdict':<20}")
    print("-" * 56)
    for r in results:
        new_str = "YES" if r["is_new"] else "NO"
        print(f"{r['name']:<12} {new_str:<12} {r['distance']:<12} {r['verdict']:<20}")

    mutations_clustered = sum(1 for r in results[1:] if not r["is_new"])
    total_mutations = len(results) - 1
    print()
    print(f"DDoD Mitigation: {mutations_clustered}/{total_mutations} mutations correctly clustered")

    if mutations_clustered == total_mutations:
        print("PASS")
        sys.exit(0)
    else:
        print("FAIL")
        sys.exit(1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("binary", help="Path to the PE binary to test")
    parser.add_argument("--api-url", default="http://localhost:8000")
    parser.add_argument("--mutations", type=int, default=10)
    args = parser.parse_args()
    run_validation(args.binary, args.api_url, args.mutations)