tests: fix and reorganize

pe/superpe.py

@@ -332,7 +332,7 @@ class SuperPe():
         iat = {}
         for entry in self.pe.DIRECTORY_ENTRY_IMPORT:
             for imp in entry.imports:
-                dll_name = entry.dll.decode('utf-8')
+                dll_name = entry.dll.decode('utf-8').lower()
                 if imp.name == None:
                     continue
                 imp_name = imp.name.decode('utf-8')
@@ -344,25 +344,17 @@ class SuperPe():
         return iat
     
 
-    def get_iat_name_for(self, dll_name: str, func_name: str) -> str:
+    def get_replacement_iat_for(self, dll_name: str, func_name: str) -> str:
+        dll_name = dll_name.lower()
         iat = self.get_iat_entries()
+        if not dll_name in iat:
+            raise Exception("DLL not found in IAT")
+
         for entry in iat[dll_name]:
             if len(entry.func_name) >= len(func_name):
                 return entry.func_name
         return None
-    
 
-    def get_iat_offset_by_nr(self, dll_name: str, nr: int) -> int:
-        encoded_dllname = dll_name
-
-        for entry in self.pe.DIRECTORY_ENTRY_IMPORT:
-            dllname = entry.dll.decode("ascii").rstrip("\x00").lower()
-            if dllname != encoded_dllname:
-                continue
-            
-            return entry.imports[nr].name_offset
-        return None
-    
 
     def get_iat_offset_by_name(self, dll_name: str, func_name: str) -> int:
         # Iterate over the imported modules and their imported functions
@@ -400,9 +392,6 @@ class SuperPe():
             offset, func_name, new_name_bytes.decode()))
         self.pe.set_bytes_at_offset(offset, new_name_bytes)
 
-        #res = self.get_iat_offset_by_name(dll_name, new_func_name)
-        #logger.info("-> RES: {}".format(res))
-
 
     ## Helpers
 

phases/injector.py

@@ -47,15 +47,9 @@ def inject_exe(
     # Patch IAT if necessary
     if source_style == FunctionInvokeStyle.iat_reuse:
         for iatRequest in project.carrier.get_all_iat_requests():
-            iat_name = superpe.get_iat_name_for("KERNEL32.dll", iatRequest.name)
+            iat_name = superpe.get_replacement_iat_for("KERNEL32.dll", iatRequest.name)
             superpe.patch_iat_entry("KERNEL32.dll", iat_name, iatRequest.name)
 
-        #iat_name_a = superpe.get_iat_name_for("KERNEL32.dll", "GetEnvironmentVariableW")
-        #iat_name_b = superpe.get_iat_name_for("KERNEL32.dll", "VirtualProtect")
-        #logger.info("Using: {} and {}".format(iat_name_a, iat_name_b))
-
-        #superpe.patch_iat_entry("KERNEL32.dll", iat_name_a, "GetEnvironmentVariableW")
-        #superpe.patch_iat_entry("KERNEL32.dll", iat_name_b, "VirtualProtect")
         superpe.pe.parse_data_directories()
 
     shellcode_offset: int = 0  # file offset

tests/data/data_reuse_pre_fixup.asm.test

@@ -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.test

@@ -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/test_asm.py

@@ -1,77 +0,0 @@
-import shutil
-from typing import List
-import unittest
-import logging
-
-from phases.asmparser import parse_asm_file
-from model.exehost import ExeHost
-from model.defs import *
-from model.carrier import Carrier
-from observer import observer
-
-
-class AsmTest(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls):
-        observer.active = False
-        
-    def test_asm_fixup(self):
-        path_in: FilePath = "tests/data/peb_walk_pre_fixup.asm"
-        path_working: FilePath = "tests/data/peb_walk_pre_fixup.asm.test"
-        carrier = Carrier()
-
-        shutil.copy(path_in, path_working)
-        parse_asm_file(carrier, path_working)
-        with open(path_working, "r") as f:
-            lines = f.readlines()
-
-        # cmp     DWORD PTR n$1[rsp], 11223344            ; 00ab4130H
-        # cmp     DWORD PTR n$1[rsp], 272         ; 00ab4130H
-        #self.assertTrue(", 272" in lines[192-1])
-        #self.assertTrue("11223344" not in lines[192-1])
-        
-        # mov     r8, QWORD PTR supermega_payload
-        # lea	r8, [shcstart]
-        self.assertTrue("lea	r8, [shcstart]" in lines[198-1])
-        self.assertTrue("supermega_payload" not in lines[198-1])
-
-        # shcstart:
-        self.assertTrue("shcstart:" in lines[212-1])
-
-        os.remove(path_working)
-
-
-    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)
-
-        carrier = Carrier()
-        parse_asm_file(carrier, path_working)
-
-        self.assertEqual(len(carrier.iat_requests), 2)
-
-        req1 = carrier.iat_requests[0]
-        self.assertEqual(req1.name, "GetEnvironmentVariableW")
-        self.assertTrue(len(req1.placeholder), 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()
-        
-        # added ; at the beginning
-        #self.assertTrue(lines[13-1].startswith("; EXTRN	__imp_GetEnvironmentVariableW:PROC"))
-        #self.assertTrue(lines[14-1].startswith("; EXTRN	__imp_VirtualAlloc:PROC"))
-
-        # 	call	QWORD PTR __imp_GetEnvironmentVariableW
-        # 	DB 044H, 0aeH, 06cH, 0b6H, 072H, 07cH
-        self.assertTrue(lines[124-1].startswith("	DB "))
-
-        # 	call	QWORD PTR __imp_VirtualAlloc
-        # 	DB 0c7H, 0b6H, 0feH, 0dcH, 0b2H, 0c6H
-        self.assertTrue(lines[148-1].startswith("	DB "))
-
-        os.remove(path_working)

tests/test_asmparser.py

@@ -0,0 +1,103 @@
+from typing import List
+import unittest
+import logging
+
+from model.defs import *
+from model.carrier import Carrier
+from observer import observer
+from helper import *
+from phases.asmparser import parse_asm_file
+from phases.masmshc import masm_shc
+
+
+def print_lines(data):
+    for i, line in enumerate(data):
+        print(f"{i+1:3}: {line}")
+
+
+class AsmTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        observer.active = False
+
+        
+    def test_asm_fixup(self):
+        asm_in: FilePath = "tests/data/peb_walk_pre_fixup.asm"
+        asm_text = file_readall_text(asm_in)
+        carrier = Carrier("fake.exe")
+        asm_text_lines = parse_asm_file(carrier, asm_text)
+
+        # cmp     DWORD PTR n$1[rsp], 11223344            ; 00ab4130H
+        # cmp     DWORD PTR n$1[rsp], 272         ; 00ab4130H
+        #self.assertTrue(", 272" in lines[192-1])
+        #self.assertTrue("11223344" not in lines[192-1])
+        
+        # mov     r8, QWORD PTR supermega_payload
+        # lea	r8, [shcstart]
+        self.assertTrue("lea	r8, [shcstart]" in asm_text_lines[198-1])
+        self.assertTrue("supermega_payload" not in asm_text_lines[198-1])
+
+        # shcstart:
+        self.assertTrue("shcstart:" in asm_text_lines[213-1])
+
+
+    def test_asm_iat_request(self):
+        asm_in: FilePath = "tests/data/iat_reuse_pre_fixup.asm"
+        asm_text = file_readall_text(asm_in)
+        carrier = Carrier("fake.exe")
+        asm_text_lines = parse_asm_file(carrier, asm_text)
+
+        self.assertEqual(len(carrier.iat_requests), 2)
+
+        req1 = carrier.iat_requests[0]
+        self.assertEqual(req1.name, "GetEnvironmentVariableW")
+        self.assertTrue(len(req1.placeholder), 6) # 6 random bytes
+        
+        req2 = carrier.iat_requests[1]
+        self.assertEqual(req2.name, "VirtualAlloc")
+        self.assertTrue(len(req2.placeholder), 6) # 6 random bytes
+
+        # added ; at the beginning
+        #self.assertTrue(lines[13-1].startswith("; EXTRN	__imp_GetEnvironmentVariableW:PROC"))
+        #self.assertTrue(lines[14-1].startswith("; EXTRN	__imp_VirtualAlloc:PROC"))
+
+        # 	call	QWORD PTR __imp_GetEnvironmentVariableW
+        # 	DB 044H, 0aeH, 06cH, 0b6H, 072H, 07cH
+        self.assertTrue(asm_text_lines[124-1].startswith("	DB "))
+
+        # 	call	QWORD PTR __imp_VirtualAlloc
+        # 	DB 0c7H, 0b6H, 0feH, 0dcH, 0b2H, 0c6H
+        self.assertTrue(asm_text_lines[148-1].startswith("	DB "))
+
+
+    def test_data_reuse_entries(self):
+        asm_in = "tests/data/data_reuse_pre_fixup.asm"
+        asm_text = file_readall_text(asm_in)
+        carrier = Carrier("fake.exe")
+        asm_text_lines = parse_asm_file(carrier, asm_text)
+        asm_text = masm_shc(asm_text_lines)  # optional here
+
+        data_reuse_entries = carrier.get_all_reusedata_fixups()
+        self.assertEqual(2, len(data_reuse_entries))
+
+        entry = data_reuse_entries[0]
+        self.assertTrue('$SG72513' in entry.string_ref)
+        self.assertTrue('rcx' in entry.register)
+        self.assertEqual(entry.data, b"U\x00S\x00E\x00R\x00P\x00R\x00O\x00F\x00I\x00L\x00E\x00\x00\x00")
+        self.assertEqual(entry.addr, 0)
+        self.assertEqual(7, len(entry.randbytes))  # needs to be 7!
+
+        entry = data_reuse_entries[1]
+        self.assertTrue('$SG72514' in entry.string_ref)
+
+
+    def test_data_reuse_fixup(self):
+        asm_in = "tests/data/data_reuse_pre_fixup.asm"
+        asm_text = file_readall_text(asm_in)
+
+        carrier = Carrier("fake.exe")
+        asm_text_lines = parse_asm_file(carrier, asm_text)
+
+        self.assertTrue("\tDB " in asm_text_lines[108-1])
+        self.assertFalse("OFFSET FLAT:$SG" in asm_text_lines[108-1])
+        

tests/test_datareuse.py

@@ -3,18 +3,35 @@ from typing import List
 import unittest
 import logging
 import os
+
 from model.defs import *
-from model.exehost import ExeHost
+
-from model.carrier import Carrier
+from pe.superpe import SuperPe
-from phases.asmparser import parse_asm_file
+from model.rangemanager import RangeManager
+from helper import *
 
 
 class DataReuseTest(unittest.TestCase):
-    def test_relocation_list(self):
+    def test_rangemanager(self):
-        exe_host = ExeHost(PATH_EXES + "7z.exe")
+        """Test RangeManager for basic functionality"""
-        exe_host.init()
+        rm = RangeManager(0, 100)
+        rm.add_range(0, 10)
+        rm.add_range(20, 30)
+        rm.add_range(50, 60)
 
-        relocs = exe_host.get_relocations_for_section(".rdata")
+        hole = rm.find_hole(10)
+        self.assertEqual((11, 19), hole)
+
+        holes = rm.find_holes(20)
+        self.assertEqual([(31, 49), (61, 100)], holes)
+
+        largest = rm.find_largest_gap()
+        self.assertEqual(40, largest)
+
+
+    def test_relocation_list(self):
+        superpe = SuperPe(PATH_EXES + "7z.exe")
+        relocs = superpe.get_relocations_for_section(".rdata")
         self.assertEqual(842, len(relocs))
         reloc = relocs[0]
         self.assertEqual(393216, reloc.base_rva)
@@ -23,10 +40,19 @@ class DataReuseTest(unittest.TestCase):
         self.assertEqual("I", reloc.type)
 
 
+    def test_relocmanager(self):
+        """Test reference EXE reloc manager information"""
+        superpe = SuperPe(PATH_EXES + "procexp64.exe")
+        rm = superpe.get_rdata_relocmanager()
+        self.assertEqual(69, len(rm.intervals))
+        # 0x1ab0 is magic currently (should use find_first_utf16_string_offset()
+        hole = rm.find_hole(20)
+        self.assertEqual(hole, (1167361, 1173015))
+
+
     def test_largestgap(self):
-        exe_host = ExeHost(PATH_EXES + "7z.exe")
+        superpe = SuperPe(PATH_EXES + "7z.exe")
-        exe_host.init()
+        rm = superpe.get_rdata_relocmanager()
-        rm = exe_host.get_rdata_relocmanager()
         start, stop = rm.find_hole(100)
         self.assertEqual(393233, start)
         self.assertEqual(394295, stop)
@@ -39,42 +65,3 @@ class DataReuseTest(unittest.TestCase):
     def test_asm_lea_create(self):
         pass
 
-
-    def test_data_reuse_entries(self):
-        asm_in = "tests/data/data_reuse_pre_fixup.asm"
-        asm_working = "tests/data/data_reuse_pre_fixup.asm.test"
-        
-        shutil.copy(asm_in, asm_working)
-        carrier = Carrier()
-        parse_asm_file(carrier, asm_working)
-        data_reuse_entries = carrier.get_all_reusedata_fixups()
-
-        self.assertEqual(2, len(data_reuse_entries))
-
-        entry = data_reuse_entries[0]
-        self.assertTrue('$SG72513' in entry.string_ref)
-        self.assertTrue('rcx' in entry.register)
-        self.assertEqual(entry.data, b"U\x00S\x00E\x00R\x00P\x00R\x00O\x00F\x00I\x00L\x00E\x00\x00\x00")
-        self.assertEqual(entry.addr, 0)
-        self.assertEqual(7, len(entry.randbytes))  # needs to be 7!
-
-        entry = data_reuse_entries[1]
-        self.assertTrue('$SG72514' in entry.string_ref)
-
-        os.remove(asm_working)
-
-
-    def test_data_reuse_fixup(self):
-        asm_in = "tests/data/data_reuse_pre_fixup.asm"
-        asm_working = asm_in + ".test"
-        
-        shutil.copy(asm_in, asm_working)
-        carrier = Carrier()
-        parse_asm_file(carrier, asm_working)
-
-        with open(asm_working, "r") as f:
-            lines = f.readlines()
-        self.assertTrue("\tDB " in lines[108-1])
-        self.assertFalse("OFFSET FLAT:$SG" in lines[108-1])
-        
-        os.remove(asm_working)

tests/test_rangemanager.py

@@ -1,37 +0,0 @@
-from typing import List
-import unittest
-import logging
-
-from model.defs import *
-from pe.superpe import SuperPe
-from model.exehost import ExeHost
-from model.rangemanager import RangeManager
-
-
-class RangeManagerTest(unittest.TestCase):
-    def test_rangemanager(self):
-        rm = RangeManager(0, 100)
-        rm.add_range(0, 10)
-        rm.add_range(20, 30)
-        rm.add_range(50, 60)
-
-        hole = rm.find_hole(10)
-        self.assertEqual((11, 19), hole)
-
-        holes = rm.find_holes(20)
-        self.assertEqual([(31, 49), (61, 100)], holes)
-
-        largest = rm.find_largest_gap()
-        self.assertEqual(40, largest)
-
-
-    def test_relocmanager(self):
-        exehost = ExeHost(PATH_EXES + "procexp64.exe")
-        exehost.init()
-        section = exehost.superpe.get_section_by_name(".rdata")
-        rm = exehost.get_rdata_relocmanager()
-        self.assertEqual(69, len(rm.intervals))
-        # 0x1ab0 is magic currently (should use find_first_utf16_string_offset()
-        #rm.add_range(section.virt_addr, section.virt_addr + 0x1AB0)
-        hole = rm.find_hole(20)
-        self.assertEqual(hole, (1167361, 1173015))

tests/test_superpe.py

@@ -9,7 +9,62 @@ from model.defs import *
 class SuperPeTest(unittest.TestCase):
 
     def test_exe(self):
-        pass
+        dll_filepath = PATH_EXES + "procexp64.exe"
+        superpe = SuperPe(dll_filepath)
+
+        # Properties
+        self.assertFalse(superpe.is_dll())
+        self.assertTrue(superpe.is_64())
+        self.assertFalse(superpe.is_dotnet())
+        self.assertEqual(superpe.get_entrypoint(), 0xE1D78)
+        self.assertIsNone(superpe.get_rwx_section())
+
+        self.assertEqual(superpe.get_image_base(), 0x140000000)
+        self.assertEqual(superpe.is_dynamic_base(), True)
+
+        # Text Section 1 (pefile SectionStructure)
+        code_sect: pefile.SectionStructure = superpe.get_code_section()
+        self.assertEqual(code_sect.Name.decode(), ".text\x00\x00\x00")
+        self.assertEqual(code_sect.VirtualAddress, 0x1000)
+        self.assertEqual(code_sect.Misc_VirtualSize, 0x11B0CE)
+
+        # Text Section 2 (PeSection)
+        code_pesect: PeSection = superpe.get_section_by_name(".text")
+        self.assertEqual(code_pesect.name, ".text")
+        self.assertEqual(code_pesect.virt_addr, 0x1000)
+        self.assertEqual(code_pesect.virt_size, 0x11B0CE)
+
+        # Relocations
+        base_relocs: List[PeRelocEntry] = superpe.get_base_relocs()
+        self.assertEqual(len(base_relocs), 2888)
+        base_reloc = base_relocs[0]
+        self.assertEqual(base_reloc.rva, 0x11E618)
+        self.assertEqual(base_reloc.base_rva, 0x11E000)
+        self.assertEqual(base_reloc.offset, 0x618)
+
+        # IAT
+        iat_entries: Dict[str, IatEntry] = superpe.get_iat_entries()
+        self.assertEqual(len(iat_entries), 24)
+        self.assertTrue("kernel32.dll" in iat_entries)
+        self.assertTrue("uxtheme.dll" in iat_entries)
+        kernel32_entries = iat_entries["kernel32.dll"]
+        self.assertEqual(len(kernel32_entries), 218)
+        entry = kernel32_entries[0]
+        self.assertEqual(entry.dll_name, "kernel32.dll")
+        self.assertEqual(entry.func_name, "FileTimeToLocalFileTime")
+        self.assertEqual(entry.iat_vaddr, 0x14011D528)
+
+        self.assertEqual(superpe.get_vaddr_of_iatentry("FileTimeToLocalFileTime"), 0x14011D528)
+        self.assertEqual(superpe.get_replacement_iat_for(
+            "kernel32.dll", "GetEnvironmentStringsW"), "FileTimeToLocalFileTime")
+
+        # Exports
+        exports = superpe.get_exports_full()
+        self.assertEqual(len(exports), 0)
+
+        # VRA/Virt to Phys/Raw
+        #raw = superpe.get_offset_from_rva(0xD690)
+        #self.assertEqual(raw, 0xCA90)
 
 
     def test_dll(self):
@@ -23,6 +78,9 @@ class SuperPeTest(unittest.TestCase):
         self.assertEqual(superpe.get_entrypoint(), 0x1350)
         self.assertIsNone(superpe.get_rwx_section())
 
+        self.assertEqual(superpe.get_image_base(), 0x1F13C0000)
+        self.assertEqual(superpe.is_dynamic_base(), True)
+
         # Text Section 1 (pefile SectionStructure)
         code_sect: pefile.SectionStructure = superpe.get_code_section()
         self.assertEqual(code_sect.Name.decode(), ".text\x00\x00\x00")
@@ -46,15 +104,19 @@ class SuperPeTest(unittest.TestCase):
         # IAT
         iat_entries: Dict[str, IatEntry] = superpe.get_iat_entries()
         self.assertEqual(len(iat_entries), 2)
-        self.assertTrue("KERNEL32.dll" in iat_entries)
+        self.assertTrue("kernel32.dll" in iat_entries)
         self.assertTrue("msvcrt.dll" in iat_entries)
-        kernel32_entries = iat_entries["KERNEL32.dll"]
+        kernel32_entries = iat_entries["kernel32.dll"]
         self.assertEqual(len(kernel32_entries), 12)
         entry = kernel32_entries[0]
-        self.assertEqual(entry.dll_name, "KERNEL32.dll")
+        self.assertEqual(entry.dll_name, "kernel32.dll")
         self.assertEqual(entry.func_name, "DeleteCriticalSection")
         self.assertEqual(entry.iat_vaddr, 0x1f13db1c4)
 
+        self.assertEqual(superpe.get_vaddr_of_iatentry("DeleteCriticalSection"), 0x1F13DB1C4)
+        self.assertEqual(superpe.get_replacement_iat_for(
+            "kernel32.dll", "GetEnvironmentStringsW"), "InitializeCriticalSection")
+
         # Exports
         exports = superpe.get_exports_full()
         self.assertEqual(len(exports), 35)