From 064b22f99e47c5b1b01c81c56b6821978a954564 Mon Sep 17 00:00:00 2001 From: hakril Date: Mon, 10 Feb 2025 10:01:55 +0100 Subject: [PATCH] Still improving tests for ARM64 --- tests/pfwtest.py | 4 ++++ tests/test_cpuid.py | 5 +++++ tests/test_native_utils.py | 14 ++++++++------ tests/test_process.py | 2 ++ tests/test_syswow.py | 2 ++ 5 files changed, 21 insertions(+), 6 deletions(-) diff --git a/tests/pfwtest.py b/tests/pfwtest.py index 923a02d..2e3969b 100644 --- a/tests/pfwtest.py +++ b/tests/pfwtest.py @@ -24,6 +24,10 @@ process_64bit_only = pytest.mark.skipif(not is_process_64_bits, reason="Test for process_syswow_only = pytest.mark.skipif(not is_process_syswow, reason="Test for syswow process only") require_admin = pytest.mark.skipif(not is_admin, reason="Test must be launched as admin") +def process_architecture_only(target_archi): + return pytest.mark.skipif(windows.current_process.architecture != target_archi, + reason="Test for {0} architecture process only".format(target_archi)) + check_for_gc_garbage = pytest.mark.usefixtures("check_for_gc_garbage") check_for_handle_leak = pytest.mark.usefixtures("check_for_handle_leak") diff --git a/tests/test_cpuid.py b/tests/test_cpuid.py index 004fe95..105ec2e 100644 --- a/tests/test_cpuid.py +++ b/tests/test_cpuid.py @@ -1,6 +1,11 @@ +import pytest + import windows +import windows.generated_def as gdef import windows.native_exec.cpuid def test_native_exec_cpuid(): + if windows.current_process.architecture == gdef.IMAGE_FILE_MACHINE_ARM64: + pytest.skip("CPUID not testable on ARM64") assert windows.native_exec.cpuid.do_cpuid(0) assert windows.native_exec.cpuid.get_proc_family_model() \ No newline at end of file diff --git a/tests/test_native_utils.py b/tests/test_native_utils.py index d20a275..9f31f5f 100644 --- a/tests/test_native_utils.py +++ b/tests/test_native_utils.py @@ -8,21 +8,23 @@ from windows.pycompat import basestring, int_types from .pfwtest import * + + @check_for_gc_garbage class TestNativeUtils(object): - @process_64bit_only + @process_architecture_only(gdef.IMAGE_FILE_MACHINE_AMD64) def test_strlenw64(self): strlenw64 = windows.native_exec.create_function(nativeutils.StrlenW64.get_code(), [gdef.UINT, gdef.LPCWSTR]) assert strlenw64("YOLO") == 4 assert strlenw64("") == 0 - @process_64bit_only + @process_architecture_only(gdef.IMAGE_FILE_MACHINE_AMD64) def test_strlena64(self): strlena64 = windows.native_exec.create_function(nativeutils.StrlenA64.get_code(), [gdef.UINT, gdef.LPCSTR]) assert strlena64(b"YOLO") == 4 assert strlena64(b"") == 0 - @process_64bit_only + @process_architecture_only(gdef.IMAGE_FILE_MACHINE_AMD64) def test_getprocaddr64(self): getprocaddr64 = windows.native_exec.create_function(nativeutils.GetProcAddress64.get_code(), [gdef.ULONG64, gdef.LPCWSTR, gdef.LPCSTR]) k32 = [mod for mod in windows.current_process.peb.modules if mod.name == "kernel32.dll"][0] @@ -37,19 +39,19 @@ class TestNativeUtils(object): assert getprocaddr64("YOLO.DLL", b"whatever") == 0xfffffffffffffffe assert getprocaddr64("KERNEL32.DLL", b"YOLOAPI") == 0xffffffffffffffff - @process_32bit_only + @process_architecture_only(gdef.IMAGE_FILE_MACHINE_I386) def test_strlenw32(self): strlenw32 = windows.native_exec.create_function(nativeutils.StrlenW32.get_code(), [gdef.UINT, gdef.LPCWSTR]) assert strlenw32("YOLO") == 4 assert strlenw32("") == 0 - @process_32bit_only + @process_architecture_only(gdef.IMAGE_FILE_MACHINE_I386) def test_strlena32(self): strlena32 = windows.native_exec.create_function(nativeutils.StrlenA32.get_code(), [gdef.UINT, gdef.LPCSTR]) assert strlena32(b"YOLO") == 4 assert strlena32(b"") == 0 - @process_32bit_only + @process_architecture_only(gdef.IMAGE_FILE_MACHINE_I386) def test_getprocaddr32(self): getprocaddr32 = windows.native_exec.create_function(nativeutils.GetProcAddress32.get_code(), [gdef.UINT, gdef.LPCWSTR, gdef.LPCSTR]) k32 = [mod for mod in windows.current_process.peb.modules if mod.name == "kernel32.dll"][0] diff --git a/tests/test_process.py b/tests/test_process.py index 2552f9a..db8b001 100644 --- a/tests/test_process.py +++ b/tests/test_process.py @@ -195,6 +195,7 @@ class TestProcessWithCheckGarbage(object): @windows_64bit_only def test_execute_to_64(self, proc64): + assert proc64.architecture == gdef.IMAGE_FILE_MACHINE_AMD64, "TODO: better machine fixture for ARM64" with proc64.allocated_memory(0x1000) as addr: shellcode = x64.MultipleInstr() shellcode += x64.Mov('RAX', 0x4242424243434343) @@ -351,6 +352,7 @@ class TestProcessWithCheckGarbage(object): @windows_64bit_only def test_set_thread_context_64(self, proc64): + assert proc64.architecture == gdef.IMAGE_FILE_MACHINE_AMD64, "TODO: better machine fixture for ARM64" code = x64.MultipleInstr() code += x64.Label(":LOOP") code += x64.Jmp(":LOOP") diff --git a/tests/test_syswow.py b/tests/test_syswow.py index 5d214d8..4a321dd 100644 --- a/tests/test_syswow.py +++ b/tests/test_syswow.py @@ -69,6 +69,8 @@ class TestSyswowRemoteProcess(object): windows.current_process.write_qword({0}, res) """.format(addr) + # Execute the import safely so that the test will not hang if import fails + proc32.execute_python("import windows") t = proc32.execute_python_unsafe(textwrap.dedent(remote_python_code)) # Wait for python execution while proc32.read_qword(addr) != 0x8877665544332211: