diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..e9e6a80 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,11 @@ +{ + "python.testing.unittestArgs": [ + "-v", + "-s", + "./tests", + "-p", + "test_*.py" + ], + "python.testing.pytestEnabled": false, + "python.testing.unittestEnabled": true +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 56270e0..a673753 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +pytest pefile capstone keystone-engine diff --git a/tests/test_asm.py b/tests/test_asm.py new file mode 100644 index 0000000..00cd841 --- /dev/null +++ b/tests/test_asm.py @@ -0,0 +1,67 @@ +import shutil +from typing import List +import unittest +import logging + +from phases.compiler import fixup_asm_file, fixup_iat_reuse +from model import ExeInfo +from defs import * + + +class AsmTest(unittest.TestCase): + 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" + + shutil.copy(path_in, path_working) + fixup_asm_file(path_working, 272) + 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[213-1]) + + os.remove(path_working) + + + def test_asm_iat_fixup(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_info = ExeInfo() + + fixup_iat_reuse(path_working, exe_info) + self.assertTrue(len(exe_info.iat_resolves), 2) + + self.assertTrue("GetEnvironmentVariableW" in exe_info.iat_resolves) + self.assertEqual(exe_info.iat_resolves["GetEnvironmentVariableW"].name, "GetEnvironmentVariableW") + self.assertEqual(exe_info.iat_resolves["GetEnvironmentVariableW"].addr, 0) + self.assertTrue(len(exe_info.iat_resolves["GetEnvironmentVariableW"].id), 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[158-1].startswith(" DB ")) + + # call QWORD PTR __imp_VirtualAlloc + # DB 0c7H, 0b6H, 0feH, 0dcH, 0b2H, 0c6H + self.assertTrue(lines[183-1].startswith(" DB ")) + + os.remove(path_working) \ No newline at end of file