From 61f7cf0a5bc120373847bd2aba9f44a4baf96abb Mon Sep 17 00:00:00 2001 From: Dobin Date: Wed, 8 May 2024 11:44:37 +0100 Subject: [PATCH] tests: fix and reorganize --- pe/superpe.py | 23 +-- phases/injector.py | 8 +- tests/data/data_reuse_pre_fixup.asm.test | 190 +++++++++++++++++++++++ tests/data/iat_reuse_pre_fixup.asm.test | 190 +++++++++++++++++++++++ tests/test_asm.py | 77 --------- tests/test_asmparser.py | 103 ++++++++++++ tests/test_datareuse.py | 85 +++++----- tests/test_rangemanager.py | 37 ----- tests/test_superpe.py | 70 ++++++++- 9 files changed, 592 insertions(+), 191 deletions(-) create mode 100644 tests/data/data_reuse_pre_fixup.asm.test create mode 100644 tests/data/iat_reuse_pre_fixup.asm.test delete mode 100644 tests/test_asm.py create mode 100644 tests/test_asmparser.py delete mode 100644 tests/test_rangemanager.py diff --git a/pe/superpe.py b/pe/superpe.py index e95c624..e8ca603 100644 --- a/pe/superpe.py +++ b/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 diff --git a/phases/injector.py b/phases/injector.py index 8242654..81e9e80 100644 --- a/phases/injector.py +++ b/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 diff --git a/tests/data/data_reuse_pre_fixup.asm.test b/tests/data/data_reuse_pre_fixup.asm.test new file mode 100644 index 0000000..26d18b3 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/data/iat_reuse_pre_fixup.asm.test b/tests/data/iat_reuse_pre_fixup.asm.test new file mode 100644 index 0000000..8040227 --- /dev/null +++ b/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 diff --git a/tests/test_asm.py b/tests/test_asm.py deleted file mode 100644 index 6133546..0000000 --- a/tests/test_asm.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/tests/test_asmparser.py b/tests/test_asmparser.py new file mode 100644 index 0000000..3d6716b --- /dev/null +++ b/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]) + diff --git a/tests/test_datareuse.py b/tests/test_datareuse.py index eb4d6a9..37c4db4 100644 --- a/tests/test_datareuse.py +++ b/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 phases.asmparser import parse_asm_file + +from pe.superpe import SuperPe +from model.rangemanager import RangeManager +from helper import * class DataReuseTest(unittest.TestCase): - def test_relocation_list(self): - exe_host = ExeHost(PATH_EXES + "7z.exe") - exe_host.init() + def test_rangemanager(self): + """Test RangeManager for basic functionality""" + 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") - exe_host.init() - rm = exe_host.get_rdata_relocmanager() + superpe = SuperPe(PATH_EXES + "7z.exe") + rm = superpe.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) diff --git a/tests/test_rangemanager.py b/tests/test_rangemanager.py deleted file mode 100644 index 5cf9429..0000000 --- a/tests/test_rangemanager.py +++ /dev/null @@ -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)) diff --git a/tests/test_superpe.py b/tests/test_superpe.py index 0c23ef3..70b22a5 100644 --- a/tests/test_superpe.py +++ b/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)