Merge pull request #12 from shaymargolis/add-different-arches-option

Add different arches option

Author: shaymargolis

.github/workflows/python-app.yml

@@ -33,4 +33,5 @@ jobs:
         flake8 . --count --max-complexity=10 --max-line-length=127 --statistics
     - name: Test with pytest
       run: |
-        pytest
+        pytest --compiler-arch=mipsbe
+        pytest --compiler-arch=mipsle

shellblocks/compiler_arch.py

@@ -0,0 +1,9 @@
+class CompilerArch():
+    def __init__(self):
+        pass
+
+    def compile_primitive(self, src_path: str) -> [str]:
+        raise NotImplementedError()
+
+    def compile_step(self, src_paths: [str], base_address: int) -> [str]:
+        raise NotImplementedError()

shellblocks/compiler_archs/__init__.py

@@ -0,0 +1,22 @@
+from enum import Enum
+from shellblocks.compiler_archs.mips import CompilerArchMIPSBE, CompilerArchMIPSLE
+from shellblocks.compiler_arch import CompilerArch
+
+
+class CompilerArchOption(Enum):
+    MIPSBE = "mipsbe"
+    MIPSLE = "mipsle"
+
+
+def compiler_arch_to_object(arch: CompilerArchOption) -> CompilerArch:
+    if arch == CompilerArchOption.MIPSBE:
+        return CompilerArchMIPSBE()
+    elif arch == CompilerArchOption.MIPSLE:
+        return CompilerArchMIPSLE()
+
+    raise NotImplementedError()
+
+
+__all__ = [
+    CompilerArchMIPSBE, CompilerArchMIPSLE,
+]

shellblocks/compiler_archs/gcc.py

@@ -0,0 +1,41 @@
+from shellblocks.compiler_arch import CompilerArch
+
+
+class CompilerArchGCC(CompilerArch):
+    def __init__(self):
+        super().__init__()
+
+        self.compiler_path = self.get_compiler_path()
+
+    def get_compiler_path(self):
+        raise NotImplementedError()
+
+    def get_ldscript_path(self):
+        raise NotImplementedError()
+
+    def get_gcc_flags(self):
+        return [
+            "-nostdlib",
+            "-ffreestanding",
+        ]
+
+    def compile_primitive(self, src_path: str) -> [str]:
+        return [
+            self.compiler_path,
+            *self.get_gcc_flags(),
+            "-c", src_path,
+            "-o", "final.o",
+            "-O3"
+        ]
+
+    def compile_step(self, src_paths: [str], base_address: int) -> [str]:
+        ldscript_loc = self.get_ldscript_path()
+
+        return [
+            self.compiler_path,
+            *src_paths,
+            "-o", "shellcode.elf",
+            *self.get_gcc_flags(),
+            f"-Wl,--section-start=.text={hex(base_address)}",
+            f"-Wl,-T{ldscript_loc}"
+        ]

shellblocks/compiler_archs/mips.py

@@ -0,0 +1,32 @@
+from shellblocks.compiler_archs.gcc import CompilerArchGCC
+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",
+        ]
+
+    def get_ldscript_path(self):
+        return (sources_location / "shellcode_ldscript.ld").as_posix()
+
+
+class CompilerArchMIPSBE(CompilerArchMIPS):
+    def get_compiler_path(self):
+        return "mips-linux-gnu-gcc-9"
+
+
+class CompilerArchMIPSLE(CompilerArchMIPS):
+    def get_gcc_flags(self):
+        return super().get_gcc_flags() + [
+            "-EL",
+        ]
+
+    def get_compiler_path(self):
+        return "mips-linux-gnu-gcc-9"

shellblocks/shellcode_primitive.py

@@ -1,6 +1,7 @@
 import shutil
 from pathlib import Path
 
+from shellblocks.compiler_arch import CompilerArch
 from shellblocks.utils import check_call_print, sources_location
 
 
@@ -37,7 +38,7 @@ def generate_header_file(self, path: Path):
     def header_requirements(self):
         return {}
 
-    def generate(self, path: Path):
+    def generate(self, path: Path, compiler: CompilerArch):
         for source in self.sources:
             source_src = sources_location / source
             source_dst = path / source
@@ -46,15 +47,10 @@ def generate(self, path: Path):
 
         self.generate_header_file(path)
 
-        check_call_print([
-            "mips-linux-gnu-gcc-9",
-            "-nostdlib",
-            "-ffreestanding",
-            "-mno-shared",
-            "-c", self.sources[0],
-            "-o", "final.o",
-            "-O3"
-        ], cwd=path.as_posix())
+        check_call_print(
+            compiler.compile_primitive(self.sources[0]),
+            cwd=path.as_posix()
+        )
 
         check_call_print([
             "objcopy",

shellblocks/shellcode_step.py

@@ -2,8 +2,9 @@
 from pathlib import Path
 
 
+from shellblocks.compiler_arch import CompilerArch
 from shellblocks.shellcode_primitive import ShellcodePrimitive
-from shellblocks.utils import check_call_print, sources_location
+from shellblocks.utils import check_call_print
 
 
 class ShellcodeStep:
@@ -13,7 +14,7 @@ def __init__(self, nickname: str, base_address: int, primitives: [ShellcodePrimi
         self.primitives = primitives
         self.max_len = max_len
 
-    def generate(self, build_dir: Path):
+    def generate(self, build_dir: Path, compiler: CompilerArch):
         # Create build dir
         try:
             shutil.rmtree(build_dir.as_posix())
@@ -29,20 +30,14 @@ def generate(self, build_dir: Path):
             p_build_dir = build_dir / p.nickname
             p_build_dir.mkdir()
 
-            out_file = p.generate(p_build_dir)
+            out_file = p.generate(p_build_dir, compiler)
             out_files.append(out_file.as_posix())
 
         # Join all primitives to final shellcode
-        ldscript_loc = (sources_location / "shellcode_ldscript.ld").as_posix()
-        check_call_print([
-            "mips-linux-gnu-gcc-9",
-            *out_files,
-            "-o", "shellcode.elf",
-            "-nostdlib",
-            "-ffreestanding",
-            f"-Wl,--section-start=.text={hex(self.base_address)}",
-            f"-Wl,-T{ldscript_loc}"
-        ], cwd=build_dir.as_posix())
+        check_call_print(
+            compiler.compile_step(out_files, self.base_address),
+            cwd=build_dir.as_posix()
+        )
 
         check_call_print([
             "objcopy",

shellblocks/test_arch_helper.py

@@ -0,0 +1,41 @@
+from shellblocks.compiler_archs import CompilerArchOption
+
+
+class ArchHelper:
+    def __init__(self, compiler_arch_option):
+        self.compiler_arch_option = compiler_arch_option
+
+
+class MIPSHelper(ArchHelper):
+    def __init__(self, compiler_arch_option):
+        super().__init__(compiler_arch_option)
+
+        assert compiler_arch_option in [
+            CompilerArchOption.MIPSBE,
+            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
+
+        if CompilerArchOption.MIPSBE == self.compiler_arch_option:
+            return val.to_bytes(4, 'big')
+        elif CompilerArchOption.MIPSLE == self.compiler_arch_option:
+            return val.to_bytes(4, 'little')
+        else:
+            raise NotImplementedError()

tests/conftest.py

@@ -4,6 +4,62 @@
 
 from pathlib import Path
 
+from unicorn import (
+    Uc,
+    UC_ARCH_MIPS,
+    UC_MODE_32,
+    UC_MODE_BIG_ENDIAN,
+    UC_MODE_LITTLE_ENDIAN,
+)
+
+
+from shellblocks.test_arch_helper import MIPSHelper
+
+from shellblocks.compiler_archs import CompilerArchOption, compiler_arch_to_object
+
+
+def pytest_addoption(parser):
+    choices = [e.value for e in CompilerArchOption]
+
+    parser.addoption(
+        "--compiler-arch",
+        type=str,
+        choices=choices,
+        default=CompilerArchOption.MIPSBE.value,
+        help="The architecture to compile to"
+    )
+
+
+@pytest.fixture
+def compiler_arch_option(request):
+    return CompilerArchOption(request.config.getoption("--compiler-arch"))
+
+
+@pytest.fixture
+def compiler_arch(compiler_arch_option):
+    return compiler_arch_to_object(compiler_arch_option)
+
+
+@pytest.fixture
+def arch_helper(compiler_arch_option):
+    if compiler_arch_option in [CompilerArchOption.MIPSBE, CompilerArchOption.MIPSLE]:
+        return MIPSHelper(compiler_arch_option)
+    else:
+        raise NotImplementedError("Arch unimplemented error!")
+
+
+@pytest.fixture
+def get_mu(compiler_arch_option):
+    def get_mu_instance():
+        if CompilerArchOption.MIPSBE == compiler_arch_option:
+            return Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_BIG_ENDIAN)
+        elif CompilerArchOption.MIPSLE == compiler_arch_option:
+            return Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_LITTLE_ENDIAN)
+        else:
+            raise NotImplementedError("Arch unimplemented error!")
+
+    return get_mu_instance
+
 
 @pytest.fixture(scope='function')
 def temp_dir_path():

tests/test_goto.py

@@ -1,6 +1,5 @@
 import pytest
 
-from unicorn import Uc, UC_ARCH_MIPS, UC_MODE_32, UC_MODE_BIG_ENDIAN
 from unicorn.mips_const import UC_MIPS_REG_PC
 
 from shellblocks.shellcode_step import ShellcodeStep
@@ -16,7 +15,7 @@
     (0xbc000000, 0xbcf00010),
     (0x91000000, 0x91000118),
 ])
-def test_goto_sanity(temp_dir_path, goto_page_and_address):
+def test_goto_sanity(get_mu, temp_dir_path, compiler_arch, goto_page_and_address):
     # Generate shellcode
     # ------------------
     shellcode_address = 0xbfc00000
@@ -31,13 +30,13 @@ def test_goto_sanity(temp_dir_path, goto_page_and_address):
         0x1000
     )
 
-    out_file = step.generate(temp_dir_path / step.nickname)
+    out_file = step.generate(temp_dir_path / step.nickname, compiler_arch)
     shellcode = out_file.read_bytes()
 
     # Try to run shellcode
     # --------------------
 
-    mu = Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_BIG_ENDIAN)
+    mu = get_mu()
     mu.mem_map(shellcode_address, 0x2000)
 
     # write machine code to be emulated to memory
@@ -54,7 +53,7 @@ def test_goto_sanity(temp_dir_path, goto_page_and_address):
     (0xbcf00010),
     (0x91000118),
 ])
-def test_goto_is_pic(temp_dir_path, shellcode_run_addr):
+def test_goto_is_pic(get_mu, temp_dir_path, compiler_arch, shellcode_run_addr):
     # Generate shellcode
     # ------------------
     shellcode_address = 0xbfc00000
@@ -71,13 +70,13 @@ def test_goto_is_pic(temp_dir_path, shellcode_run_addr):
         0x1000
     )
 
-    out_file = step.generate(temp_dir_path / step.nickname)
+    out_file = step.generate(temp_dir_path / step.nickname, compiler_arch)
     shellcode = out_file.read_bytes()
 
     # Try to run shellcode
     # --------------------
 
-    mu = Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_BIG_ENDIAN)
+    mu = get_mu()
     mu.mem_map(shellcode_run_sector, 0x2000)
 
     # write machine code to be emulated to memory