From 7d2a5f31b81569190865ba0bda355ea504610ca4 Mon Sep 17 00:00:00 2001 From: dannyp303 Date: Tue, 12 Mar 2024 15:30:52 -0400 Subject: [PATCH] add lief add/remove section modifier (#443) * add lief add/remove section modifier * Changelog * typo --------- Co-authored-by: Dan Pesce --- ofrak_core/CHANGELOG.md | 1 + ofrak_core/ofrak/core/elf/lief_modifier.py | 51 +++++++++++++++++++ .../components/test_elf_modifiers.py | 42 ++++++++++++++- 3 files changed, 93 insertions(+), 1 deletion(-) diff --git a/ofrak_core/CHANGELOG.md b/ofrak_core/CHANGELOG.md index 7c76dda21..71a472a7e 100644 --- a/ofrak_core/CHANGELOG.md +++ b/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)) diff --git a/ofrak_core/ofrak/core/elf/lief_modifier.py b/ofrak_core/ofrak/core/elf/lief_modifier.py index b15c40e64..a6774eed3 100644 --- a/ofrak_core/ofrak/core/elf/lief_modifier.py +++ b/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) diff --git a/ofrak_core/test_ofrak/components/test_elf_modifiers.py b/ofrak_core/test_ofrak/components/test_elf_modifiers.py index d3714fae9..b302e45cc 100644 --- a/ofrak_core/test_ofrak/components/test_elf_modifiers.py +++ b/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