refactor: DataReuseEntry(), IatEntry()

model/carrier.py

@@ -1,7 +1,6 @@
 from typing import Dict, List
 import logging
 
-from model.exehost import DataReuseEntry
 
 logger = logging.getLogger("Carrier")
 
@@ -12,6 +11,14 @@ class IatRequest():
         self.placeholder: bytes = placeholder    # Random bytes as placeholder
 
 
+class DataReuseEntry():
+    def __init__(self, string_ref: str, register: str, randbytes: bytes):
+        self.string_ref = string_ref
+        self.register = register
+        self.randbytes = randbytes
+        self.data = b''
+        self.addr = 0
+
 
 class Carrier():
     def __init__(self):

model/exehost.py

@@ -18,19 +18,10 @@ class RelocEntry():
 
 
 class IatEntry():
-    def __init__(self, dll_name, func_name, func_addr):
+    def __init__(self, dll_name: str, func_name: str, iat_vaddr: int):
-        self.dll_name = dll_name
+        self.dll_name: str = dll_name
-        self.func_name = func_name
+        self.func_name: str = func_name
-        self.func_addr = func_addr
+        self.iat_vaddr: int = iat_vaddr
-
-
-class DataReuseEntry():
-    def __init__(self, string_ref: str, register: str, randbytes: bytes):
-        self.string_ref = string_ref
-        self.register = register
-        self.randbytes = randbytes
-        self.data = b''
-        self.addr = 0
 
 
 class ExeHost():
@@ -42,7 +33,7 @@ class ExeHost():
         self.pe: pefile.PE = None
         self.superpe: SuperPe = None
 
-        self.iat = {}  # Dict[str, List[Dict[str, str]]]
+        self.iat: Dict[str, IatEntry] = {}
         self.base_relocs = []
 
         self.image_base: int = 0
@@ -128,26 +119,13 @@ class ExeHost():
                 if not dll_name in self.iat:
                     self.iat[dll_name] = []
 
-                self.iat[dll_name].append({
+                self.iat[dll_name].append(IatEntry(dll_name, imp_name, imp_addr))
-                    "dll_name": dll_name,
-                    "func_name": imp_name,
-                    "iat_vaddr": imp_addr
-                })
         
 
     def get_vaddr_of_iatentry(self, func_name: str) -> int:
         for dll_name in self.iat:
             for entry in self.iat[dll_name]:
-                if entry["func_name"] == func_name:
+                if entry.func_name == func_name:
-                    return entry["iat_vaddr"]
+                    return entry.iat_vaddr
         return None
     
-    ## Other
-
-    def print(self):
-        logger.info("--( Required IAT Resolves: ")
-        for _, cap in self.iat_requests.items():
-            if cap.addr == 0:
-                logger.info("   {:28} {}".format(cap.name, "N/A"))
-            else:
-                logger.info("   {:28} 0x{:x}".format(cap.name, cap.addr))

phases/datareuse.py

@@ -4,7 +4,7 @@ from intervaltree import Interval, IntervalTree
 from typing import List, Dict
 import os
 
-from model.exehost import DataReuseEntry
+from model.carrier import DataReuseEntry
 
 
 def bytes_to_asm_db(byte_data: bytes) -> bytes:

phases/injector.py

@@ -6,7 +6,7 @@ import time
 import tempfile
 import logging
 
-from model.carrier import Carrier
+from model.carrier import Carrier, DataReuseEntry
 from peparser.pehelper import *
 from model.exehost import *
 from observer import observer