Still improving tests for ARM64

Author: hakril

tests/pfwtest.py

@@ -24,6 +24,10 @@
 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")

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()

tests/test_native_utils.py

@@ -8,21 +8,23 @@
 
 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 @@ def test_getprocaddr64(self):
         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]

tests/test_process.py

@@ -195,6 +195,7 @@ def test_execute_to_proc32(self, proc32):
 
     @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 @@ def test_set_thread_context_32(self, proc32):
 
     @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")

tests/test_syswow.py

@@ -69,6 +69,8 @@ def test_getset_syswow_context(self, proc32):
         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: