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