add lief add/remove section modifier (redballoonsecurity#443)

* add lief add/remove section modifier

* Changelog

* typo

---------

Co-authored-by: Dan Pesce <[email protected]>

ofrak_core/CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## [Unreleased](https://github.com/redballoonsecurity/ofrak/tree/master)
 ### Added
+- Add modifier to add and remove sections using lief. ([#443](https://github.com/redballoonsecurity/ofrak/pull/443))
 - Add tabbed content views and a decompilation view to the OFRAK GUI. ([#436](https://github.com/redballoonsecurity/ofrak/pull/436/))
 - Refactor HexView and related components to use mousewheel instead of scroll and compartmentalize all comonents to src/hex. ([#427](https://github.com/redballoonsecurity/ofrak/pull/427))
 - Add an improved ISO9660 packer that leverages `mkisofs` instead of PyCdLib. ([#393](https://github.com/redballoonsecurity/ofrak/pull/393))

ofrak_core/ofrak/core/elf/lief_modifier.py

@@ -73,3 +73,54 @@ async def modify(self, resource: Resource, config: LiefAddSegmentConfig) -> None
  new_data = f_handle.read()
  # replace all old content (old range) with new content from Lief
  resource.queue_patch(Range(0, await resource.get_data_length()), new_data)
+
+
+@dataclass
+class LiefAddSectionModifierConfig(ComponentConfig):
+ name: str
+ content: bytes
+ flags: int
+
+
+class LiefAddSectionModifer(Modifier[LiefAddSectionModifierConfig]):
+ targets = (Elf,)
+
+ async def modify(self, resource: Resource, config: LiefAddSectionModifierConfig):
+ binary: lief.ELF.Binary = lief.parse(await resource.get_data())
+ section: lief.ELF.Section = lief.ELF.Section()
+ section.name = config.name
+ section.content = list(config.content)
+ section.flags = config.flags
+ binary.add(section)
+
+ with tempfile.NamedTemporaryFile() as temp_file:
+ binary.write(temp_file.name)
+ temp_file.flush()
+ with open(temp_file.name, "rb") as f_handle:
+ new_data = f_handle.read()
+ # replace all old content (old range) with new content from Lief
+ resource.queue_patch(Range(0, await resource.get_data_length()), new_data)
+
+
+@dataclass
+class LiefRemoveSectionModifierConfig(ComponentConfig):
+ name: str
+
+
+class LiefRemoveSectionModifier(Modifier[LiefRemoveSectionModifierConfig]):
+ targets = (Elf,)
+
+ async def modify(self, resource: Resource, config: LiefRemoveSectionModifierConfig):
+ binary: lief.ELF.Binary = lief.parse(await resource.get_data())
+ section: lief.ELF.Section = binary.get_section(config.name)
+ if section is None:
+ raise AttributeError(f"No section with name {config.name}")
+ binary.remove(section)
+
+ with tempfile.NamedTemporaryFile() as temp_file:
+ binary.write(temp_file.name)
+ temp_file.flush()
+ with open(temp_file.name, "rb") as f_handle:
+ new_data = f_handle.read()
+ # replace all old content (old range) with new content from Lief
+ resource.queue_patch(Range(0, await resource.get_data_length()), new_data)

ofrak_core/test_ofrak/components/test_elf_modifiers.py

@@ -1,11 +1,19 @@
 import os
 import subprocess
+from typing import Optional
 
 import pytest
 from elftools.elf.elffile import ELFFile
 from test_ofrak.components.hello_world_elf import hello_elf
 
-from ofrak.core import LiefAddSegmentConfig, LiefAddSegmentModifier
+from ofrak.core import (
+ LiefAddSegmentConfig,
+ LiefAddSegmentModifier,
+ LiefAddSectionModifer,
+ LiefAddSectionModifierConfig,
+ LiefRemoveSectionModifier,
+ LiefRemoveSectionModifierConfig,
+)
 from ofrak.service.resource_service_i import ResourceFilter
 from ofrak.core.elf.model import (
  Elf,
@@ -254,3 +262,35 @@ def assert_segment_exists(filepath: str, vaddr: int, length: int):
  if segment.header.p_vaddr == vaddr and segment.header.p_memsz == length:
  return
  raise ValueError("Could not find segment in binary")
+
+
+async def test_lief_add_section_modifier(hello_out: Resource, tmp_path):
+ config = LiefAddSectionModifierConfig(name=".test", content=b"test", flags=0)
+ await hello_out.run(LiefAddSectionModifer, config=config)
+ elf_path = tmp_path / "test.elf"
+ await hello_out.flush_data_to_disk(elf_path)
+ assert segment_exists(elf_path, ".test", content=b"test")
+
+
+async def test_lief_remove_section_modifier(hello_out: Resource, tmp_path):
+ original = tmp_path / "original.elf"
+ await hello_out.flush_data_to_disk(original)
+ assert segment_exists(original, ".text")
+ config = LiefRemoveSectionModifierConfig(name=".text")
+ await hello_out.run(LiefRemoveSectionModifier, config=config)
+ modified = tmp_path / "modified.elf"
+ await hello_out.flush_data_to_disk(modified)
+ assert not segment_exists(modified, ".text")
+
+
+def segment_exists(filepath: str, name: str, content: Optional[bytes] = None):
+ with open(filepath, "rb") as f:
+ elffile = ELFFile(f)
+ sections = list(elffile.iter_sections())
+ for section in sections:
+ if section.name == name:
+ if content is not None and content in section.data():
+ return True
+ if content is None:
+ return True
+ return False