Merge pull request #21 from shaymargolis/add-lots-of-archs

shaymargolis · web-flow · commit c26012d62578 · 2025-02-16T10:01:54.000+02:00
Add lots of archs
diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
@@ -36,3 +36,5 @@ jobs:
         pytest --compiler-arch=mipsbe
         pytest --compiler-arch=mipsle
         pytest --compiler-arch=armle
+        pytest --compiler-arch=x86
+        pytest --compiler-arch=x86_64
diff --git a/shellblocks/compiler_arch.py b/shellblocks/compiler_arch.py
@@ -1,10 +1,14 @@
 from typing import List
+from shellblocks.compiler_arch_option import CompilerArchOption
 
 
 class CompilerArch():
     def __init__(self):
         pass
 
+    def compiler_arch_option(self) -> [CompilerArchOption]:
+        raise NotImplementedError()
+
     def compile_primitive(self, src_path: str) -> List[str]:
         raise NotImplementedError()
 
diff --git a/shellblocks/compiler_arch_option.py b/shellblocks/compiler_arch_option.py
@@ -0,0 +1,32 @@
+from subprocess import check_output, CalledProcessError
+from enum import Enum
+
+
+class CompilerArchOption(Enum):
+    MIPSBE = "mipsbe"
+    MIPSLE = "mipsle"
+    ARMLE = "armle"
+    X86 = "x86"
+    X86_64 = "x86_64"
+
+
+def get_current_platform():
+    try:
+        machine_bytes = check_output(["gcc", "-dumpmachine"])
+    except FileNotFoundError:
+        return None
+    except CalledProcessError:
+        return None
+
+    machine = machine_bytes.decode()
+
+    if "mips-" in machine:
+        return CompilerArchOption.MIPSBE
+    if "mipsel-" in machine:
+        return CompilerArchOption.MIPSLE
+    if "arm-" in machine:
+        return CompilerArchOption.ARMLE
+    if "x86_64-" in machine:
+        return CompilerArchOption.X86_64
+
+    return None
diff --git a/shellblocks/compiler_archs/__init__.py b/shellblocks/compiler_archs/__init__.py
@@ -1,22 +1,22 @@
-from enum import Enum
+from shellblocks.compiler_arch_option import CompilerArchOption
 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"
-
-
 def compiler_arch_to_object(arch: CompilerArchOption) -> CompilerArch:
     if arch == CompilerArchOption.MIPSBE:
         return CompilerArchMIPSBE()
     elif arch == CompilerArchOption.MIPSLE:
         return CompilerArchMIPSLE()
     elif arch == CompilerArchOption.ARMLE:
         return CompilerArchARMLE()
+    elif arch == CompilerArchOption.X86:
+        return CompilerArchX86()
+    elif arch == CompilerArchOption.X86_64:
+        return CompilerArchX86_64()
 
     raise NotImplementedError()
 
diff --git a/shellblocks/compiler_archs/arm.py b/shellblocks/compiler_archs/arm.py
@@ -1,5 +1,6 @@
 from typing import List
 
+from shellblocks.compiler_arch_option import CompilerArchOption
 from shellblocks.compiler_archs.gcc import CompilerArchGCC
 from shellblocks.utils import sources_location
 
@@ -8,8 +9,6 @@ class CompilerArchARM(CompilerArchGCC):
     def __init__(self):
         super().__init__()
 
-        self.compiler_path = self.get_compiler_path()
-
     def get_headers(self) -> List[str]:
         return ["arch/arm/utils.h"]
 
@@ -20,3 +19,6 @@ def get_ldscript_path(self):
 class CompilerArchARMLE(CompilerArchARM):
     def get_compiler_path(self):
         return "arm-linux-gnueabi-gcc-10"
+
+    def compiler_arch_option(self) -> [CompilerArchOption]:
+        return [CompilerArchOption.MIPSLE]
diff --git a/shellblocks/compiler_archs/gcc.py b/shellblocks/compiler_archs/gcc.py
@@ -1,11 +1,15 @@
 from shellblocks.compiler_arch import CompilerArch
+from shellblocks.compiler_arch_option import get_current_platform
 
 
 class CompilerArchGCC(CompilerArch):
     def __init__(self):
         super().__init__()
 
-        self.compiler_path = self.get_compiler_path()
+        if get_current_platform() in self.compiler_arch_option():
+            self.compiler_path = "gcc"
+        else:
+            self.compiler_path = self.get_compiler_path()
 
     def get_compiler_path(self):
         raise NotImplementedError()
@@ -15,8 +19,14 @@ def get_ldscript_path(self):
 
     def get_gcc_flags(self):
         return [
+            # Dont try to link with any std library
+            # or assume hosted assumptions: Allow any entrypoint
+            # and function definition
             "-nostdlib",
             "-ffreestanding",
+            # Require GCC to keep the compilation order
+            # as declared in the source files
+            "-fno-toplevel-reorder",
         ]
 
     def compile_primitive(self, src_path: str) -> [str]:
diff --git a/shellblocks/compiler_archs/mips.py b/shellblocks/compiler_archs/mips.py
@@ -1,15 +1,14 @@
 from typing import List
 
 from shellblocks.compiler_archs.gcc import CompilerArchGCC
+from shellblocks.compiler_arch_option import CompilerArchOption
 from shellblocks.utils import sources_location
 
 
 class CompilerArchMIPS(CompilerArchGCC):
     def __init__(self):
         super().__init__()
 
-        self.compiler_path = self.get_compiler_path()
-
     def get_gcc_flags(self):
         return super().get_gcc_flags() + [
             "-mno-shared",
@@ -26,6 +25,9 @@ class CompilerArchMIPSBE(CompilerArchMIPS):
     def get_compiler_path(self):
         return "mips-linux-gnu-gcc-9"
 
+    def compiler_arch_option(self) -> [CompilerArchOption]:
+        return [CompilerArchOption.MIPSBE]
+
 
 class CompilerArchMIPSLE(CompilerArchMIPS):
     def get_gcc_flags(self):
@@ -35,3 +37,6 @@ def get_gcc_flags(self):
 
     def get_compiler_path(self):
         return "mips-linux-gnu-gcc-9"
+
+    def compiler_arch_option(self) -> [CompilerArchOption]:
+        return [CompilerArchOption.MIPSLE]
diff --git a/shellblocks/compiler_archs/x86.py b/shellblocks/compiler_archs/x86.py
@@ -0,0 +1,15 @@
+from shellblocks.compiler_arch_option import CompilerArchOption
+from shellblocks.compiler_archs.x86_64 import CompilerArchX86_64
+
+
+class CompilerArchX86(CompilerArchX86_64):
+    def __init__(self):
+        super().__init__()
+
+    def get_gcc_flags(self):
+        return super().get_gcc_flags() + [
+            "-m32",
+        ]
+
+    def compiler_arch_option(self) -> [CompilerArchOption]:
+        return [CompilerArchOption.X86, CompilerArchOption.X86_64]
diff --git a/shellblocks/compiler_archs/x86_64.py b/shellblocks/compiler_archs/x86_64.py
@@ -0,0 +1,22 @@
+from typing import List
+
+from shellblocks.compiler_arch_option import CompilerArchOption
+from shellblocks.compiler_archs.gcc import CompilerArchGCC
+from shellblocks.utils import sources_location
+
+
+class CompilerArchX86_64(CompilerArchGCC):
+    def __init__(self):
+        super().__init__()
+
+    def get_headers(self) -> List[str]:
+        return ["arch/x86/utils.h"]
+
+    def get_ldscript_path(self):
+        return (sources_location / "shellcode_ldscript.ld").as_posix()
+
+    def get_compiler_path(self):
+        return "x86_64-linux-gnu-gcc-8"
+
+    def compiler_arch_option(self) -> [CompilerArchOption]:
+        return [CompilerArchOption.X86_64]
diff --git a/shellblocks/src/arch/mips/utils.h b/shellblocks/src/arch/mips/utils.h
@@ -9,8 +9,7 @@
         "nop\n\t"                                               \
         "bal " GET_REL_LABEL(LABEL) "\n\t"                      \
         "nop\n\t"                                               \
-        GET_REL_LABEL(LABEL) ": move $v0, $ra\n\t"              \
-        "move %0, $v0\n\t"                                      \
+        GET_REL_LABEL(LABEL) ": move %0, $ra\n\t"               \
         "addiu %0, (" #LABEL " - " GET_REL_LABEL(LABEL) ")\n\t" \
         : "=r" (OUTPUT) : : "$ra"                               \
     )
diff --git a/shellblocks/src/arch/utils.h b/shellblocks/src/arch/utils.h
@@ -3,6 +3,8 @@
 
 #if defined(__mips__)
 #include "mips/utils.h"
+#elif defined(__i386__) || defined(__x86_64__)
+#include "x86/utils.h"
 #elif defined(__arm__)
 #include "arm/utils.h"
 #endif
diff --git a/shellblocks/src/arch/x86/utils.h b/shellblocks/src/arch/x86/utils.h
@@ -0,0 +1,14 @@
+#ifndef SHELLCODE_BLOCKS_ARCH_I386_UTILS_H
+#define SHELLCODE_BLOCKS_ARCH_I386_UTILS_H
+
+#define NOP_OPCODE (0x90909090)
+
+#define GET_REL_ADDRESS(OUTPUT, LABEL)                         \
+    __asm__ volatile (                                         \
+        "call " GET_REL_LABEL(LABEL) "\n\t"                    \
+        GET_REL_LABEL(LABEL) ": pop %0\n\t"                    \
+        "add $(" #LABEL " - " GET_REL_LABEL(LABEL) "), %0\n\t" \
+        : "=r" (OUTPUT) : :                                    \
+    )
+
+#endif // !SHELLCODE_BLOCKS_ARCH_I386_UTILS_H
diff --git a/shellblocks/src/memset.c b/shellblocks/src/memset.c
@@ -2,11 +2,6 @@
 #include "arch/utils.h"
 #include "memset.h"
 
-// Must be the first variable
-DECLARE_NOP_END(nop_end);
-
-const u8 items[] REL_ACCESS_STRING = MEMSET_ITEMS;
-
 void start(void) {
     u8 *items_rel;
 
@@ -23,3 +18,8 @@ void start(void) {
 
     JUMP_TO_NOP_END(nop_end);
 }
+
+const u8 items[] REL_ACCESS_STRING = MEMSET_ITEMS;
+
+// Must be the last variable
+DECLARE_NOP_END(nop_end);
diff --git a/shellblocks/src/print.c b/shellblocks/src/print.c
@@ -2,11 +2,6 @@
 #include "arch/utils.h"
 #include "print.h"
 
-// Must be the first variable
-DECLARE_NOP_END(nop_end);
-
-const char print_string[] REL_ACCESS_STRING = PRINT_STRING;
-
 void start(void) {
     void (*print_func)(const char *) = (void (*)(const char *))PRINT_FUNCTION_ADDRESS;
     char *print_string_rel;
@@ -16,3 +11,8 @@ void start(void) {
 
     JUMP_TO_NOP_END(nop_end);
 }
+
+const char print_string[] REL_ACCESS_STRING = PRINT_STRING;
+
+// Must be the last variable
+DECLARE_NOP_END(nop_end);
diff --git a/shellblocks/test_arch_helper.py b/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,54 @@ 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_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)
@@ -17,9 +66,6 @@ def __init__(self, compiler_arch_option):
             CompilerArchOption.ARMLE,
         ]
 
-    def get_jump_hook_bytes(self, jump_hook_goto):
-        raise NotImplementedError()
-
     def get_ret_bytes(self):
         val = 0xe12fff1e  # "bx lr" in ARM
 
@@ -53,21 +99,6 @@ def __init__(self, compiler_arch_option):
             CompilerArchOption.MIPSLE,
         ]
 
-    def get_jump_hook_bytes(self, jump_hook_goto):
-        EXPECTED_HOOK = [
-            0x3c020000 + (jump_hook_goto >> 16),
-            0x24420000 + (jump_hook_goto & 0xffff),
-            0x00400008,
-            0x00000000,
-        ]
-
-        if CompilerArchOption.MIPSBE == self.compiler_arch_option:
-            return b"".join(map(lambda x: x.to_bytes(4, 'big'), EXPECTED_HOOK))
-        elif CompilerArchOption.MIPSLE == self.compiler_arch_option:
-            return b"".join(map(lambda x: x.to_bytes(4, 'little'), EXPECTED_HOOK))
-        else:
-            raise NotImplementedError()
-
     def get_ret_bytes(self):
         val = 0x03e00008  # "jr $ra" in MIPS
 
diff --git a/tests/conftest.py b/tests/conftest.py
