Simple class for loading a PE file in Unicorn

pemulator.py

from typing import List

from capstone import *
from capstone.x86 import *
from unicorn import *
from unicorn.x86_const import *
from pefile import PE


class Emulator():
    """
    Base class for emulators.
    """

    emu: Uc
    image_base: int
    entry_point: int
    stack_base: int
    stack_size: int

    _pe: PE
    _cs: Cs
    _memory_alignment: int
    _mode: int
    _debug: bool
    _trace: bool

    def __init__(
        self,
        pe_data: bytes,
        image_base: int = None,
        stack_base: int = None,
        stack_size: int = None,
        mode: int = 32,
        debug: bool = False,
        trace: bool = False,
    ):
        """
        Initializes the Emulator instance.
        """
        self._debug = debug
        self._set_mode(mode)
        # Load pe file
        self._pe = PE(data=pe_data)
        
        # Get the base address of the module
        if image_base is None:
            self.image_base = self._pe.OPTIONAL_HEADER.ImageBase
        else:
            self.image_base = image_base

        # Get the entry point of the module
        self.entry_point = self._pe.OPTIONAL_HEADER.AddressOfEntryPoint
        
        # Get virtual alignment 
        self._memory_alignment = self._pe.OPTIONAL_HEADER.SectionAlignment

        # Load PE into emulator
        self.add_data_sections()

        # Get lower bounds of allocated memory from emulator
        memory_base = self.get_memory_lower_bound()

        # Get upper bounds of allocated memory from emulator
        memory_top = self.get_memory_upper_bound()
        
        # Get the stack base
        if stack_base is None:
            self.stack_base = self.memory_align(0x1000)
        else:
            self.stack_base = self.memory_align(stack_base)
        
        # Get the stack size
        if stack_size is None:
            self.stack_size = self.memory_align(0x5000)
        else:
            self.stack_size = self.memory_align(stack_size)
        
        # Adjust the stack base and size if needed
        if self.stack_base +  self.stack_size > memory_base and self.stack_base < memory_top:
            self.stack_base = self.memory_align(memory_top + 0x1000)

        # Map the stack
        self.emu.mem_map(self.stack_base, self.stack_size)
        if self._mode == 32:
            tmp_ESP = self.stack_base + self.stack_size // 2
            self.emu.reg_write(UC_X86_REG_ESP, tmp_ESP)
            self.emu.reg_write(UC_X86_REG_EBP, tmp_ESP)
        elif self._mode == 64:
            tmp_RSP = self.stack_base + self.stack_size // 2
            self.emu.reg_write(UC_X86_REG_RSP, tmp_RSP)
            self.emu.reg_write(UC_X86_REG_RBP, tmp_RSP)
        
        # Enable trace
        self._trace = trace
        if self._trace:
            self.emu.hook_add(UC_HOOK_CODE, self.trace_code)
            self._debug = True

    def debug_print(self, message: str) -> None:
        """
        Prints a debug message if debugging is enabled.
        """
        if self._debug:
            print(message)

    def _set_mode(self, mode: int) -> None:
        """
        Sets 32/64 bit mode for the emulator.
        """
        self._mode = mode
        if mode == 32:
            self.emu = Uc(UC_ARCH_X86, UC_MODE_32)
            self._cs = Cs(CS_ARCH_X86, CS_MODE_32)
            self._cs.detail = True
        elif mode == 64:
            self.emu = Uc(UC_ARCH_X86, UC_MODE_64)
            self._cs = Cs(CS_ARCH_X86, CS_MODE_64)
            self._cs.detail = True
        else:
            raise Exception("Invalid mode")
        
    def get_memory_lower_bound(self) -> int:
        """
        Returns the lower bound of the allocated memory.
        """
        memory_segments = list(self.emu.mem_regions())
        memory_segments.sort(key=lambda x: x[0])
        memory_base = 0
        if len(memory_segments) > 0:
            memory_base = memory_segments[0][0]
        return memory_base
    
    def get_memory_upper_bound(self) -> int:
        """
        Returns the upper bound of the allocated memory.
        """
        memory_segments = list(self.emu.mem_regions())
        memory_segments.sort(key=lambda x: x[1], reverse=True)
        memory_top = 0
        if len(memory_segments) > 0:
            memory_top = self.memory_align(memory_segments[0][1])
        return memory_top
    
    def memory_align(self, address: int) -> int:
        """
        Aligns the given address to the nearest multiple of alignment.
        """
        return ((address + self._memory_alignment - 1) // self._memory_alignment) * self._memory_alignment

    def add_data_sections(self) -> None:
        """
        Adds sections to emulator
        """
        # For each section in the PE file add it to the emulator
        for section in self._pe.sections:
            # Get the section data
            data = section.get_data()
            # Get the section size
            size = section.Misc_VirtualSize
            # Align the section size
            size_aligned = self.memory_align(size)
            # Get the section address
            address = self.image_base + section.VirtualAddress
            permissions = 0
            # Check if the section is readable
            if section.Characteristics & 0x40000000:
                permissions |= UC_PROT_READ
            # Check if the section is writable
            if section.Characteristics & 0x80000000:
                permissions |= UC_PROT_WRITE
            # Check if the section is executable
            if section.Characteristics & 0x20000000:
                permissions |= UC_PROT_EXEC

            # Map the memory with the combined permissions
            self.debug_print(f"Mapping section {section.Name.decode()} at 0x{address:x} with size 0x{size_aligned:x} and permissions {permissions}")
            self.emu.mem_map(address, size_aligned, permissions)
            self.emu.mem_write(address, data)

        return 
    
    def get_iat_address_from_import(self, library: str, import_name: str) -> int:
        """
        Get the address of the import address table entry for the given import name.
        """
        orig_pe_base = self._pe.OPTIONAL_HEADER.ImageBase
        for entry in self._pe.DIRECTORY_ENTRY_IMPORT:
            if entry.dll.decode().lower() == library.lower():
                for imp in entry.imports:
                    if imp.name.decode().lower() == import_name.lower():
                        return imp.address - orig_pe_base + self.image_base
        return None
    
    def trace_code(self, uc: Uc, address: int, size: int, user_data: object) -> None:
        """
        Adds an emulator hook for UC_HOOK_CODE used to print all executed instructions.
        """
        # Get the current instruction
        insn = next(self._cs.disasm(uc.mem_read(address, size), address))
        self.debug_print("0x%x:\t%s\t%s" % (insn.address, insn.mnemonic, insn.op_str))