refactor: IAT stuff to ExeHost

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():

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
 
 

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:

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")

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:

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

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

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