tests: Add x86 and x86_64 archs

Author: shaymargolis

shellblocks/compiler_archs/__init__.py

@@ -1,13 +1,17 @@
 from enum import Enum
 from shellblocks.compiler_archs.mips import CompilerArchMIPSBE, CompilerArchMIPSLE
 from shellblocks.compiler_archs.arm import CompilerArchARMLE
+from shellblocks.compiler_archs.x86 import CompilerArchX86
+from shellblocks.compiler_archs.x86_64 import CompilerArchX86_64
 from shellblocks.compiler_arch import CompilerArch
 
 
 class CompilerArchOption(Enum):
     MIPSBE = "mipsbe"
     MIPSLE = "mipsle"
     ARMLE = "armle"
+    X86 = "x86"
+    X86_64 = "x86_64"
 
 
 def compiler_arch_to_object(arch: CompilerArchOption) -> CompilerArch:
@@ -17,6 +21,10 @@ def compiler_arch_to_object(arch: CompilerArchOption) -> CompilerArch:
         return CompilerArchMIPSLE()
     elif arch == CompilerArchOption.ARMLE:
         return CompilerArchARMLE()
+    elif arch == CompilerArchOption.X86:
+        return CompilerArchX86()
+    elif arch == CompilerArchOption.X86_64:
+        return CompilerArchX86_64()
 
     raise NotImplementedError()
 

shellblocks/test_arch_helper.py

@@ -1,5 +1,6 @@
 from unicorn.mips_const import UC_MIPS_REG_PC, UC_MIPS_REG_29, UC_MIPS_REG_4
 from unicorn.arm_const import UC_ARM_REG_PC, UC_ARM_REG_SP, UC_ARM_REG_R0
+from unicorn.x86_const import UC_X86_REG_EIP, UC_X86_REG_ESP, UC_X86_REG_RDI
 
 from shellblocks.compiler_archs import CompilerArchOption
 
@@ -9,6 +10,57 @@ def __init__(self, compiler_arch_option):
         self.compiler_arch_option = compiler_arch_option
 
 
+class X86Helper(ArchHelper):
+    def __init__(self, compiler_arch_option):
+        super().__init__(compiler_arch_option)
+
+        assert compiler_arch_option in [
+            CompilerArchOption.X86,
+            CompilerArchOption.X86_64,
+        ]
+
+        self.word_size = 4
+
+        if CompilerArchOption.X86_64 == self.compiler_arch_option:
+            self.word_size = 8
+
+    def get_jump_hook_bytes(self, jump_hook_goto):
+        raise NotImplementedError()
+
+    def get_ret_bytes(self):
+        return b"\xc3"  # "ret" in x86
+
+    def get_curr_pc(self, mu):
+        return mu.reg_read(UC_X86_REG_EIP)
+
+    def set_curr_sp(self, mu, new_stack):
+        mu.reg_write(UC_X86_REG_ESP, new_stack)
+
+    def get_curr_sp(self, mu):
+        return mu.reg_read(UC_X86_REG_ESP)
+
+    def get_curr_func_arg(self, mu, func_arg):
+        if CompilerArchOption.X86_64 == self.compiler_arch_option:
+            if func_arg == 0:
+                return mu.reg_read(UC_X86_REG_RDI)
+
+            raise NotImplementedError(f"Getting {func_arg} func arg")
+
+        if CompilerArchOption.X86 == self.compiler_arch_option:
+            # cdecl
+
+            val = mu.mem_read(
+                (mu.reg_read(UC_X86_REG_ESP)
+                    + (func_arg + 1) * self.word_size), self.word_size
+            )
+
+            return int.from_bytes(val, 'little')
+
+        raise NotImplementedError(
+            "Only x86 or x86-64 calling conventions are implemented"
+        )
+
+
 class ARMHelper(ArchHelper):
     def __init__(self, compiler_arch_option):
         super().__init__(compiler_arch_option)

tests/conftest.py

@@ -8,14 +8,16 @@
     Uc,
     UC_ARCH_MIPS,
     UC_ARCH_ARM,
+    UC_ARCH_X86,
     UC_MODE_ARM,
     UC_MODE_32,
+    UC_MODE_64,
     UC_MODE_BIG_ENDIAN,
     UC_MODE_LITTLE_ENDIAN,
 )
 
 
-from shellblocks.test_arch_helper import MIPSHelper, ARMHelper
+from shellblocks.test_arch_helper import MIPSHelper, ARMHelper, X86Helper
 
 from shellblocks.compiler_archs import CompilerArchOption, compiler_arch_to_object
 
@@ -48,6 +50,8 @@ def arch_helper(compiler_arch_option):
         return MIPSHelper(compiler_arch_option)
     if compiler_arch_option in [CompilerArchOption.ARMLE]:
         return ARMHelper(compiler_arch_option)
+    if compiler_arch_option in [CompilerArchOption.X86, CompilerArchOption.X86_64]:
+        return X86Helper(compiler_arch_option)
     else:
         raise NotImplementedError("Arch unimplemented error!")
 
@@ -61,6 +65,10 @@ def get_mu_instance():
             return Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_LITTLE_ENDIAN)
         elif CompilerArchOption.ARMLE == compiler_arch_option:
             return Uc(UC_ARCH_ARM, UC_MODE_ARM | UC_MODE_LITTLE_ENDIAN)
+        elif CompilerArchOption.X86 == compiler_arch_option:
+            return Uc(UC_ARCH_X86, UC_MODE_32)
+        elif CompilerArchOption.X86_64 == compiler_arch_option:
+            return Uc(UC_ARCH_X86, UC_MODE_64)
         else:
             raise NotImplementedError("Arch unimplemented error!")