Injection fixes (#172)

* Fix bug when injecting multiple payloads * AVR fixups * Allow setting physical address in LiefAddSegmentModifier * Fix ghidra GetInstructions null pointer * Black formatting * Delete all descendants after SegmentInjectorModifier runs * Update changelogs Co-authored-by: Andrés Hernández <[email protected]>
redballoonsecurity · Jan 24, 2023 · 85a3c7a · 85a3c7a
1 parent 0d936f8
commit 85a3c7a
Show file tree

Hide file tree

Showing 7 changed files with 47 additions and 16 deletions.
diff --git a/disassemblers/ofrak_ghidra/ofrak_ghidra/ghidra_scripts/GetInstructions.java b/disassemblers/ofrak_ghidra/ofrak_ghidra/ghidra_scripts/GetInstructions.java
@@ -67,7 +67,7 @@ class Result {
             Instruction instruction = getInstructionAt(startAddr);
 
             if (instruction == null) {
-                instruction = instruction.getNext();
+                instruction = getInstructionAfter(startAddr);
             }
 
             this.instructions = new ArrayList<>();

diff --git a/ofrak_core/CHANGELOG.md b/ofrak_core/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## [2.1.0](https://github.com/redballoonsecurity/ofrak/compare/ofrak-v2.0.0...ofrak-v2.1.0) - 2023-01-20
 ### Added
+- `LiefAddSegmentConfig` now has an optional `physical_address` argument.
 - New `identify` and `unpack` subcommands to CLI [#164](https://github.com/redballoonsecurity/ofrak/pull/164)
 - Move GUI server to `ofrak_core`, startup GUI through CLI, add testing for server, make GUI pip installable. [#168](https://github.com/redballoonsecurity/ofrak/pull/168)
   - `python -m ofrak gui` starts the OFRAK GUI server.
@@ -22,6 +23,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Fixed
 - Remove unneeded and slow `.save()` when unpacking filesystems [#171](https://github.com/redballoonsecurity/ofrak/pull/171)
+- Fixed null pointer bug in Ghidra scripts.
+- `SegmentInjectorModifier` deletes all descendants of the modified section, fixing a bug that would arrise when applying more than one segment modification.
 - Added missing elf section types to ElfSectionType [#178](https://github.com/redballoonsecurity/ofrak/pull/178)
 - Handled .rela placeholder segments in the same fashion as .got [#185](https://github.com/redballoonsecurity/ofrak/pull/185)
 

diff --git a/ofrak_core/ofrak/core/elf/lief_modifier.py b/ofrak_core/ofrak/core/elf/lief_modifier.py
@@ -1,6 +1,6 @@
 import tempfile
 from dataclasses import dataclass
-from typing import List
+from typing import List, Optional
 
 import lief
 
@@ -22,13 +22,15 @@ class LiefAddSegmentConfig(ComponentConfig):
     :ivar rwx_flags: string representation of the new segment's R/W/X permissions
     :ivar replace_note: replace the unused NOTE segment with the new segment, rather than adding a
     new segment. defaults to `True`, as adding a new segment may corrupt the ELF due to a LIEF bug.
+    :ivar physical_address: overwrite the default physical address (defaults to the virtual address)
     """
 
     virtual_address: int
     alignment: int
     content: List[int]
     rwx_flags: str
     replace_note: bool = True
+    physical_address: Optional[int] = None
 
 
 class LiefAddSegmentModifier(Modifier[LiefAddSegmentConfig]):
@@ -43,6 +45,8 @@ async def modify(self, resource: Resource, config: LiefAddSegmentConfig) -> None
         segment.content = config.content
         segment.alignment = config.alignment
         segment.virtual_address = config.virtual_address
+        if config.physical_address is not None:
+            segment.physical_address = config.physical_address
         if "r" in config.rwx_flags:
             segment.add(lief.ELF.SEGMENT_FLAGS.R)
         if "w" in config.rwx_flags:
@@ -56,7 +60,9 @@ async def modify(self, resource: Resource, config: LiefAddSegmentConfig) -> None
             #   and https://github.com/lief-project/LIEF/issues/143
             if not binary.has(lief.ELF.SEGMENT_TYPES.NOTE):
                 raise ValueError("Binary must have a NOTE section to add a new section")
-            _ = binary.replace(segment, binary[lief.ELF.SEGMENT_TYPES.NOTE])
+            segment = binary.replace(segment, binary[lief.ELF.SEGMENT_TYPES.NOTE])
+            if config.physical_address is not None:
+                segment.physical_address = config.physical_address
         else:
             _ = binary.add(segment)
 

diff --git a/ofrak_core/ofrak/core/patch_maker/modifiers.py b/ofrak_core/ofrak/core/patch_maker/modifiers.py
@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import os
 import tempfile
@@ -175,7 +176,7 @@ async def modify(self, resource: Resource, config: SegmentInjectorModifierConfig
             )
         )
 
-        injection_tasks = []
+        injection_tasks: List[Tuple[Resource, BinaryInjectorModifierConfig]] = []
 
         for segment, segment_data in config.segments_and_data:
             if segment.length == 0 or segment.vm_address == 0:
@@ -208,8 +209,16 @@ async def modify(self, resource: Resource, config: SegmentInjectorModifierConfig
 
             injection_tasks.append((region.resource, BinaryInjectorModifierConfig(patches)))
 
+        all_decendants = await resource.get_descendants()
         for injected_resource, injection_config in injection_tasks:
-            await injected_resource.run(BinaryInjectorModifier, injection_config)
+            result = await injected_resource.run(BinaryInjectorModifier, injection_config)
+            parents_to_delete = list(
+                filter(lambda r: r.get_id() in result.resources_modified, all_decendants)
+            )
+            to_delete = []
+            for parent in parents_to_delete:
+                to_delete.extend(list(await parent.get_descendants()))
+            await asyncio.gather(*(r.delete() for r in to_delete))
 
 
 @dataclass

diff --git a/ofrak_patch_maker/CHANGELOG.md b/ofrak_patch_maker/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## [3.0.0](https://github.com/redballoonsecurity/ofrak/compare/ofrak-patch-maker-v.2.0.0...ofrak-patch-maker-v.3.0.0) - 2023-01-20
 ### Added
+- `-fno-optimize-sibling-calls` flag added to AVR toolchain.
 - Optional permission map parameter to `Allocatable.allocate_bom`, which enables developers to express where 
 segments of one set of permissions may be placed in the destination binary. For example, a developer may specify
 to place `MemoryPermissions.R` `Segments` in destination program `MemoryRegions` of `MemoryPermissions.R` 
@@ -17,6 +18,11 @@ or `MemoryPermissions.RX`.
 - Make toolchain names in `toolchain.conf` more specific:
   - `GNU_ARM_NONE` changed to `GNU_ARM_NONE_EABI_10_2_1`.
   - `GNU_X86_64_LINUX` changed to `GNU_X86_64_LINUX_EABI_10_3_0`.
+- Pass `-mmcu` value to the AVR preprocessor.
+- Raise a more descriptive error on toolchain failure.
+
+### Fixed
+- Toolchain `preprocess()` method now returns the path to the preprocessed file.
 
 ### Removed
 - Removed `ToolchainVersion`.

diff --git a/ofrak_patch_maker/ofrak_patch_maker/toolchain/abstract.py b/ofrak_patch_maker/ofrak_patch_maker/toolchain/abstract.py
@@ -255,15 +255,19 @@ def _execute_tool(
 
         args = [tool_path] + final_flags + in_files
         self._logger.info(" ".join(args))
-        if env:
-            my_env = os.environ.copy()
-            my_env.update(env)
-            self._logger.info(f"With env: {my_env}")
-            proc = subprocess.run(
-                args, stdout=subprocess.PIPE, encoding="utf-8", check=True, env=my_env
-            )
-        else:
-            proc = subprocess.run(args, stdout=subprocess.PIPE, encoding="utf-8", check=True)
+        try:
+            if env:
+                my_env = os.environ.copy()
+                my_env.update(env)
+                self._logger.info(f"With env: {my_env}")
+                proc = subprocess.run(
+                    args, stdout=subprocess.PIPE, encoding="utf-8", check=True, env=my_env
+                )
+            else:
+                proc = subprocess.run(args, stdout=subprocess.PIPE, encoding="utf-8", check=True)
+        except subprocess.CalledProcessError as e:
+            cmd = " ".join(args)
+            raise ValueError(f'Command "{cmd}" returned non-zero exit status {e.returncode}')
 
         return proc.stdout
 
@@ -281,8 +285,7 @@ def preprocess(self, source_file: str, header_dirs: List[str], out_dir: str = ".
             [source_file] + ["-I" + x for x in header_dirs],
             out_file=out_file,
         )
-        # Note that we return the source file here!!!
-        return os.path.abspath(source_file)
+        return os.path.abspath(out_file)
 
     def compile(self, c_file: str, header_dirs: List[str], out_dir: str = ".") -> str:
         """

diff --git a/ofrak_patch_maker/ofrak_patch_maker/toolchain/gnu_avr.py b/ofrak_patch_maker/ofrak_patch_maker/toolchain/gnu_avr.py
@@ -35,8 +35,12 @@ def __init__(
             self._assembler_flags.append(
                 f"-mmcu={self._config.assembler_cpu or processor.sub_isa.value}"
             )
+            self._preprocessor_flags.append(
+                f"-mmcu={self._config.compiler_cpu or processor.sub_isa.value}"
+            )
         else:
             raise ValueError("sub_isa is required for AVR linking")
+        self._compiler_flags.append("-fno-optimize-sibling-calls")
 
     @property
     def name(self) -> str: