#!/usr/bin/env python3
"""
validate_aes_cbc.py : valide bit-exact, par execution du vrai code machine
(Unicorn, x86 32 bits), la chaine AES-256-CBC utilisee par le payload
InterLock deballe (FUN_00402fc0 / thunk_FUN_00403210), en appelant
sequentiellement les fonctions LibTomCrypt reelles extraites du binaire :

  1. FUN_0040cb10(&PTR_DAT_00437900)               -- register_cipher("aes")
  2. FUN_0040c7b0("aes")                            -- find_cipher("aes") -> index
  3. FUN_0040d720(index, IV, key, 32, 0, &cbc_ctx)  -- cbc_start
  4. FUN_0040d6a0(IV, 16, &cbc_ctx)                 -- cbc_setiv
  5. FUN_0040d500(pt, ct, len, &cbc_ctx)            -- cbc_encrypt

Anciennement bloque a l'etape 1 : le harnais dereferencait
PTR_DAT_00437900 (lisait la valeur qu'il contient, 0x4378e0) et passait
CE POINTEUR a register_cipher, ce qui aboutissait a "aes\0" suivi de
zeros -- un `name` de type char* dont la valeur numerique n'est meme pas
une adresse valide dans l'image (verifie : hors de SizeOfImage). Erreur
d'indirection : les deux vrais appelants de register_cipher dans le
binaire (`FUN_00402a80`, et surtout `FUN_00401e10` qui est le flux complet
reel de dechiffrement AES-CBC du malware, cle/IV hex -> dechiffrement d'un
buffer global) font tous deux `FUN_0040cb10(&PTR_DAT_00437900)`, c'est a
dire l'ADRESSE DE LA VARIABLE ELLE-MEME (0x437900), pas sa valeur
dereferencee. En relisant la structure a partir de 0x437900 directement :
name=0x4378e0, ID=6, min/max/block/rounds=16/32/16/10, setup=0x40a220,
ecb_encrypt=0x408f30, ecb_decrypt=0x409310, test=0x40a2d0, done=0x408f20,
keysize=0x409790 -- exactement le descripteur AES textbook de LibTomCrypt,
deja entierement present et valide dans le fichier statique. Aucune
initialisation CRT n'est necessaire : la structure est un POD constant
resolu au link-time, jamais modifiee au runtime (aucune reference
d'ecriture trouvee vers cette adresse dans tout le binaire). Il suffisait
de corriger cette seule adresse.

Usage : validate_aes_cbc.py
"""
import struct

import pefile
from unicorn import Uc, UC_ARCH_X86, UC_MODE_32, UC_HOOK_CODE, UcError
from unicorn.x86_const import UC_X86_REG_EAX, UC_X86_REG_ECX, UC_X86_REG_EDX, UC_X86_REG_ESP, UC_X86_REG_EBP

STUB_BASE = 0x00720000
STUB_SIZE = 0x00010000

SAMPLE = "tools/unpacked/9b31237dbf87d06caa3769bec491fad2dae131846e323bed636a6c6b8e24d7ae.exe"

REGISTER_CIPHER = 0x0040cb10
FIND_CIPHER = 0x0040c7b0
CBC_START = 0x0040d720
CBC_SETIV = 0x0040d6a0
CBC_ENCRYPT = 0x0040d500
AES_DESC_PTR_LOC = 0x00437900  # contient l'adresse du descripteur AES (PTR_DAT_00437900)

STACK_BASE = 0x00700000
STACK_SIZE = 0x00010000
SCRATCH_BASE = 0x00710000
SCRATCH_SIZE = 0x00010000
RET_SENTINEL = 0x00666666

# Vecteur de test genere via pycryptodome (source de confiance), pas une
# constante recopiee de memoire (source d'erreurs de frappe repetees ici) :
from Crypto.Cipher import AES as _RefAES

KEY = bytes(range(32))          # cle de test simple, n'importe laquelle convient
IV = bytes(range(16, 32))       # IV de test, distinct de la cle
PLAINTEXT = bytes(range(16))    # un seul bloc de 16 octets
EXPECTED_CIPHERTEXT = _RefAES.new(KEY, _RefAES.MODE_CBC, IV).encrypt(PLAINTEXT)
assert len(KEY) == 32 and len(IV) == 16 and len(PLAINTEXT) == 16 and len(EXPECTED_CIPHERTEXT) == 16


def main():
    pe = pefile.PE(SAMPLE, fast_load=True)
    pe.parse_data_directories()
    image_base = pe.OPTIONAL_HEADER.ImageBase

    mu = Uc(UC_ARCH_X86, UC_MODE_32)
    size_of_headers = pe.OPTIONAL_HEADER.SizeOfHeaders
    map_size = (size_of_headers + 0xFFF) & ~0xFFF
    mu.mem_map(image_base, map_size)
    raw = bytes(pe.__data__)
    mu.mem_write(image_base, raw[:size_of_headers] + b"\x00" * (map_size - size_of_headers))

    for section in pe.sections:
        va = image_base + section.VirtualAddress
        vsize = max(section.Misc_VirtualSize, section.SizeOfRawData)
        map_size = (vsize + 0xFFF) & ~0xFFF
        if map_size == 0:
            continue
        data = section.get_data()
        mu.mem_map(va, map_size)
        towrite = data[:map_size] if len(data) >= map_size else data + b"\x00" * (map_size - len(data))
        mu.mem_write(va, towrite)

    mu.mem_map(STACK_BASE, STACK_SIZE)
    mu.mem_map(SCRATCH_BASE, SCRATCH_SIZE)
    mu.mem_map(STUB_BASE, STUB_SIZE)
    mu.mem_map(RET_SENTINEL & ~0xFFF, 0x1000)
    mu.mem_write(RET_SENTINEL, b"\xC3")

    # IAT : chaque import recoit un stub dedie (corps = RET) dans STUB_BASE,
    # avec une vraie semantique cdecl pour les quelques fonctions dont
    # find_cipher/cbc_* ont reellement besoin (strcmp, malloc, memcpy,
    # memset, memcmp) -- pas de simples no-op, le resultat doit etre correct.
    stub_by_addr = {}
    off = 0
    for entry_imp in getattr(pe, "DIRECTORY_ENTRY_IMPORT", []):
        for imp in entry_imp.imports:
            name = imp.name.decode(errors="ignore") if imp.name else f"ord{imp.ordinal}"
            stub_addr = STUB_BASE + off
            off += 0x10
            mu.mem_write(stub_addr, b"\xC3" + b"\x90" * 15)
            mu.mem_write(imp.address, struct.pack("<I", stub_addr))
            stub_by_addr[stub_addr] = name

    bump = [0]

    def bump_alloc(n):
        addr = STUB_BASE + STUB_SIZE // 2 + bump[0]
        bump[0] += (n + 0xF) & ~0xF
        return addr

    def read_cstr(uc, addr, maxlen=256):
        out = bytearray()
        for i in range(maxlen):
            b = bytes(uc.mem_read(addr + i, 1))
            if b == b"\x00":
                break
            out += b
        return bytes(out)

    def hook_stub(uc, address, size, user_data):
        name = stub_by_addr.get(address)
        if name is None:
            return
        esp = uc.reg_read(UC_X86_REG_ESP)
        args = struct.unpack("<3I", bytes(uc.mem_read(esp + 4, 12)))
        a0, a1, a2 = args
        print(f"    [stub] {name} args={hex(a0)},{hex(a1)},{hex(a2)} esp={hex(esp)}")
        if name == "strcmp":
            s1 = read_cstr(uc, a0)
            s2 = read_cstr(uc, a1)
            ret = (s1 > s2) - (s1 < s2)
            uc.reg_write(UC_X86_REG_EAX, ret & 0xFFFFFFFF)
        elif name == "memcmp":
            b1 = bytes(uc.mem_read(a0, a2))
            b2 = bytes(uc.mem_read(a1, a2))
            ret = (b1 > b2) - (b1 < b2)
            uc.reg_write(UC_X86_REG_EAX, ret & 0xFFFFFFFF)
        elif name in ("malloc", "calloc"):
            addr = bump_alloc(a0 if a0 else 0x10)
            uc.reg_write(UC_X86_REG_EAX, addr)
        elif name == "free":
            uc.reg_write(UC_X86_REG_EAX, 0)
        elif name == "memcpy":
            data = bytes(uc.mem_read(a1, a2))
            uc.mem_write(a0, data)
            uc.reg_write(UC_X86_REG_EAX, a0)
        elif name == "memset":
            uc.mem_write(a0, bytes([a1 & 0xFF]) * a2)
            uc.reg_write(UC_X86_REG_EAX, a0)
        else:
            uc.reg_write(UC_X86_REG_EAX, 0)

    mu.hook_add(UC_HOOK_CODE, hook_stub, begin=STUB_BASE, end=STUB_BASE + STUB_SIZE // 2)

    scratch_off = [0]

    def alloc(n):
        addr = SCRATCH_BASE + scratch_off[0]
        scratch_off[0] += (n + 0xF) & ~0xF
        return addr

    key_addr = alloc(32)
    mu.mem_write(key_addr, KEY)
    iv_addr = alloc(16)
    mu.mem_write(iv_addr, IV)
    pt_addr = alloc(16)
    mu.mem_write(pt_addr, PLAINTEXT)
    ct_addr = alloc(16)
    mu.mem_write(ct_addr, b"\x00" * 16)
    cbc_ctx_addr = alloc(0x1138)
    name_addr = alloc(8)
    mu.mem_write(name_addr, b"aes\x00")

    # &PTR_DAT_00437900 : l'adresse de la variable elle-meme, PAS sa valeur
    # dereferencee (c'etait le bug qui bloquait cette validation).
    aes_desc_addr = AES_DESC_PTR_LOC
    print(f"aes_desc @ {hex(aes_desc_addr)}")

    def call_cdecl(func_addr, args):
        """Appelle une fonction cdecl 32 bits : push args droite-a-gauche,
        push adresse de retour = RET_SENTINEL, saute a func_addr, tourne
        jusqu'a RET_SENTINEL, lit EAX."""
        esp = STACK_BASE + STACK_SIZE - 0x1000
        for a in reversed(args):
            esp -= 4
            mu.mem_write(esp, struct.pack("<I", a))
        esp -= 4
        mu.mem_write(esp, struct.pack("<I", RET_SENTINEL))
        mu.reg_write(UC_X86_REG_ESP, esp)
        mu.reg_write(UC_X86_REG_EBP, esp)
        mu.emu_start(func_addr, RET_SENTINEL, timeout=0, count=50_000_000)
        return mu.reg_read(UC_X86_REG_EAX)

    print("1) register_cipher(&aes_desc)")
    r1 = call_cdecl(REGISTER_CIPHER, [aes_desc_addr])
    print("   ->", r1)

    dat_00447540 = struct.unpack("<I", mu.mem_read(0x00447540, 4))[0]
    print(f"   DAT_00447540 apres register_cipher = {hex(dat_00447540)}")
    print(f"   aes_desc @ {hex(aes_desc_addr)}, contenu (32 octets): {bytes(mu.mem_read(aes_desc_addr, 32)).hex()}")

    print("2) find_cipher(\"aes\")")
    idx = call_cdecl(FIND_CIPHER, [name_addr])
    print("   -> index =", idx)

    print("3) cbc_start(idx, iv, key, 32, 0, &cbc_ctx)")
    r3 = call_cdecl(CBC_START, [idx, iv_addr, key_addr, 32, 0, cbc_ctx_addr])
    print("   ->", r3)

    print("4) cbc_setiv(iv, 16, &cbc_ctx)")
    r4 = call_cdecl(CBC_SETIV, [iv_addr, 16, cbc_ctx_addr])
    print("   ->", r4)

    print("5) cbc_encrypt(pt, ct, 16, &cbc_ctx)")
    r5 = call_cdecl(CBC_ENCRYPT, [pt_addr, ct_addr, 16, cbc_ctx_addr])
    print("   ->", r5)

    ct = bytes(mu.mem_read(ct_addr, 16))
    print("\nciphertext obtenu :", ct.hex())
    print("ciphertext attendu:", EXPECTED_CIPHERTEXT.hex())
    print("MATCH" if ct == EXPECTED_CIPHERTEXT else "MISMATCH")


if __name__ == "__main__":
    try:
        main()
    except UcError as e:
        print("UcError:", e)
        raise