File size: 3,523 Bytes
8e95c1c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | 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)
|