diff --git a/model/exehost.py b/model/exehost.py index ea6ff2c..c358ecb 100644 --- a/model/exehost.py +++ b/model/exehost.py @@ -1,9 +1,9 @@ -from typing import Dict +from typing import Dict, List import logging import pefile from model.defs import * -import pehelper +import peparser.pehelper as pehelper from peparser.misc import get_physical_address logger = logging.getLogger("ExeHost") @@ -23,10 +23,14 @@ class IatResolve(): self.id ) + class ExeHost(): def __init__(self, filepath: FilePath): self.filepath: FilePath = filepath + self.iat_resolves: Dict[str, IatResolve] = {} + self.iat = {} + self.image_base = 0 self.dynamic_base = False @@ -34,7 +38,7 @@ class ExeHost(): self.code_size = 0 self.code_section = None - self.iat = {} + self.base_relocs = [] self.rwx_section = None @@ -42,10 +46,6 @@ class ExeHost(): self.ep_raw = None - def add_iat_resolve(self, func_name, placeholder): - self.iat_resolves[func_name] = IatResolve( - func_name, placeholder, pehelper.get_addr_for(self.iat, func_name)) - def init(self): logger.info("--[ Analyzing: {}".format(self.filepath)) @@ -76,7 +76,7 @@ class ExeHost(): self.code_size)) # iat - self.iat = pehelper.extract_iat(pe) + self.iat = self.extract_iat(pe) # relocs if hasattr(pe, 'DIRECTORY_ENTRY_BASERELOC'): @@ -95,20 +95,62 @@ class ExeHost(): self.rwx_section = pehelper.get_rwx_section(pe) + ## IAT related + + def add_iat_resolve(self, func_name, placeholder): + self.iat_resolves[func_name] = IatResolve( + func_name, placeholder, self._get_addr_of_iat_function(func_name)) + + def get_all_iat_resolvs(self) -> Dict[str, IatResolve]: return self.iat_resolves - def has_all_functions(self, needs): + def has_all_iat_functions(self, needed_functions: List[str]) -> bool: is_ok = True - for func_name in needs: - addr = pehelper.get_addr_for(self.iat, func_name) + for func_name in needed_functions: + addr = self._get_addr_of_iat_function(func_name) if addr == 0: logging.info("---( Function not available as import: {}".format(func_name)) is_ok = False return is_ok + def _get_addr_of_iat_function(self, func_name: str) -> int: + for dll_name in self.iat: + for entry in self.iat[dll_name]: + if entry["func_name"] == func_name: + return entry["func_addr"] + return 0 + + def extract_iat(self, pe: pefile.PE): + iat = {} + + # If the PE file was loaded using the fast_load=True argument, we will need to parse the data directories: + #pe.parse_data_directories() + + # Retrieve the IAT entries from the PE file + for entry in pe.DIRECTORY_ENTRY_IMPORT: + for imp in entry.imports: + dll_name = entry.dll.decode('utf-8') + if imp.name == None: + continue + imp_name = imp.name.decode('utf-8') + imp_addr = imp.address + + if not dll_name in iat: + iat[dll_name] = [] + + iat[dll_name].append({ + "dll_name": dll_name, + "func_name": imp_name, + "func_addr": imp_addr + }) + + return iat + + ## Other + def print(self): logger.info("--( Required IAT Resolves: ") for _, cap in self.iat_resolves.items(): diff --git a/observer.py b/observer.py index 9df7125..22ea8f5 100644 --- a/observer.py +++ b/observer.py @@ -3,7 +3,7 @@ import pprint from capstone import Cs, CS_ARCH_X86, CS_MODE_64 from model import * -from r2helper import r2_disas +from peparser.r2helper import r2_disas from helper import delete_all_files_in_directory diff --git a/pehelper.py b/peparser/pehelper.py similarity index 77% rename from pehelper.py rename to peparser/pehelper.py index 8088fd8..cf414ab 100644 --- a/pehelper.py +++ b/peparser/pehelper.py @@ -84,43 +84,6 @@ def assemble_and_disassemble_jump(current_address: int, destination_address: int return machine_code -# IAT Stuff - -def extract_iat(pe: pefile.PE): - iat = {} - - # If the PE file was loaded using the fast_load=True argument, we will need to parse the data directories: - #pe.parse_data_directories() - - # Retrieve the IAT entries from the PE file - for entry in pe.DIRECTORY_ENTRY_IMPORT: - for imp in entry.imports: - dll_name = entry.dll.decode('utf-8') - if imp.name == None: - continue - imp_name = imp.name.decode('utf-8') - imp_addr = imp.address - - if not dll_name in iat: - iat[dll_name] = [] - - iat[dll_name].append({ - "dll_name": dll_name, - "func_name": imp_name, - "func_addr": imp_addr - }) - - return iat - - -def get_addr_for(iat, func_name: str) -> int: - for dll_name in iat: - for entry in iat[dll_name]: - if entry["func_name"] == func_name: - return entry["func_addr"] - return 0 - - ## Utils def remove_trailing_null_bytes(data: bytes) -> bytes: diff --git a/r2helper.py b/peparser/r2helper.py similarity index 100% rename from r2helper.py rename to peparser/r2helper.py diff --git a/phases/assembler.py b/phases/assembler.py index 87b17f9..d2c50ce 100644 --- a/phases/assembler.py +++ b/phases/assembler.py @@ -3,7 +3,7 @@ import logging from model import * from config import config from observer import observer -from pehelper import * +from peparser.pehelper import * from helper import * logger = logging.getLogger("Assembler") diff --git a/phases/compiler.py b/phases/compiler.py index 91dc4c4..29fe3c0 100644 --- a/phases/compiler.py +++ b/phases/compiler.py @@ -136,7 +136,7 @@ def fixup_asm_file(filename: FilePath, payload_len: int, short_call_patching: bo return True -def get_function_stubs(asm_in: FilePath): +def get_function_stubs(asm_in: FilePath) -> List[str]: functions = [] with open(asm_in, 'r', encoding='utf-8') as asmfile: diff --git a/phases/injector.py b/phases/injector.py index 49c5e71..218b8a9 100644 --- a/phases/injector.py +++ b/phases/injector.py @@ -6,7 +6,7 @@ import time import tempfile import logging -from pehelper import * +from peparser.pehelper import * from model.exehost import * from observer import observer from helper import rbrunmode_str diff --git a/supermega.py b/supermega.py index a1281f1..ac5cde0 100644 --- a/supermega.py +++ b/supermega.py @@ -16,7 +16,7 @@ import phases.compiler import phases.assembler import phases.injector from observer import observer -from pehelper import extract_code_from_exe +from peparser.pehelper import extract_code_from_exe from model.project import Project from model.settings import Settings @@ -168,7 +168,7 @@ def start(settings: Settings): # Decide if we can use IAT_REUSE (all function calls available as import) required_functions = phases.compiler.get_function_stubs(main_asm_file) - if project.exe_host.has_all_functions(required_functions): + if project.exe_host.has_all_iat_functions(required_functions): settings.source_style = SourceStyle.iat_reuse logger.warning("--[ SourceStyle: Using IAT_REUSE".format()) # all good, patch ASM diff --git a/tests/test_derbackdoorer.py b/tests/test_derbackdoorer.py index 02e5454..75484b1 100644 --- a/tests/test_derbackdoorer.py +++ b/tests/test_derbackdoorer.py @@ -5,7 +5,7 @@ import logging from model.exehost import ExeHost from model.defs import * -from pehelper import extract_code_from_exe +from peparser.pehelper import extract_code_from_exe from helper import hexdump from observer import observer