Migrate to lief

Author: insolor

dfint64_patch/binio.py

@@ -1,8 +1,6 @@
 from collections.abc import Iterable
 from typing import BinaryIO
 
-from peclasses.section_table import Section
-
 from dfint64_patch.type_aliases import Offset
 
 
@@ -40,11 +38,3 @@ def write_string(
     bs = s.encode() if encoding is None else s.encode(encoding)
 
     file_object.write(bs.ljust(new_len, b"\0"))
-
-
-def read_section_data(file: BinaryIO, section: Section) -> bytes:
-    return read_bytes(
-        file,
-        Offset(section.pointer_to_raw_data),
-        Offset(section.size_of_raw_data),
-    )

dfint64_patch/extract_strings/cli.py

@@ -1,12 +1,11 @@
 import operator
 from dataclasses import dataclass
+from io import BufferedReader
 from pathlib import Path
-from typing import BinaryIO, cast
 
+import lief
 from omegaconf import DictConfig
-from peclasses.portable_executable import PortableExecutable
 
-from dfint64_patch.binio import read_section_data
 from dfint64_patch.config import with_config
 from dfint64_patch.cross_references.cross_references_relative import (
     find_relative_cross_references,
@@ -19,25 +18,25 @@
 from dfint64_patch.utils import maybe_open
 
 
-def extract_strings(pe_file: BinaryIO) -> list[ExtractedStringInfo]:
-    pe = PortableExecutable(pe_file)
+def extract_strings(pe_file: BufferedReader) -> list[ExtractedStringInfo]:
+    pe = lief.PE.parse(pe_file)
+    assert pe is not None
 
-    sections = pe.section_table
-    code_section = sections[0]
-    string_section = sections[1]
+    code_section = pe.sections[0]
+    string_section = pe.sections[1]
 
-    image_base = pe.optional_header.image_base
+    image_base = pe.optional_header.imagebase
 
     strings = list(
         extract_strings_from_raw_bytes(
-            read_section_data(pe_file, string_section),
-            base_address=Rva(cast(int, string_section.virtual_address) + image_base),
+            string_section.content,
+            base_address=Rva(string_section.virtual_address + image_base),
         ),
     )
 
     cross_references = find_relative_cross_references(
-        read_section_data(pe_file, code_section),
-        base_address=Rva(cast(int, code_section.virtual_address) + image_base),
+        code_section.content,
+        base_address=Rva(code_section.virtual_address + image_base),
         addresses=map(operator.itemgetter(0), strings),
     )
 

dfint64_patch/extract_strings/from_raw_bytes.py

@@ -35,7 +35,7 @@ class ExtractedStringInfo(NamedTuple):
 
 
 def extract_strings_from_raw_bytes(
-    bytes_block: bytes,
+    bytes_block: bytes | memoryview,
     base_address: Rva = RVA0,
     alignment: int = 4,
     encoding: str = "cp437",
@@ -48,6 +48,7 @@ def extract_strings_from_raw_bytes(
     :param encoding: string encoding
     :return: Iterator[ExtractedStringInfo]
     """
+    bytes_block = bytes(bytes_block)
     i = 0
     while i < len(bytes_block):
         if bytes_block[i] == b"\0":

dfint64_patch/patching/patch.py

@@ -1,11 +1,9 @@
 from collections.abc import Iterable, Mapping
 from pathlib import Path
-from typing import cast
 
+import lief
 from loguru import logger
-from peclasses.portable_executable import PortableExecutable
 
-from dfint64_patch.binio import read_section_data
 from dfint64_patch.cross_references.cross_references_relative import (
     find_intersected_cross_references,
     find_relative_cross_references,
@@ -17,29 +15,27 @@
 
 def patch(patched_file: str | Path, translation_table: list[tuple[str, str]], encoding: str) -> None:
     with Path(patched_file).open("r+b") as pe_file:
-        pe = PortableExecutable(pe_file)
+        pe = lief.PE.parse(pe_file)
+        assert pe is not None
 
-        sections = pe.section_table
-        code_section = sections[0]
-        data_section = sections[1]
-
-        cast(int, pe.optional_header.image_base)
+        code_section = pe.sections[0]
+        data_section = pe.sections[1]
 
         logger.info("Extracting strings...")
         strings = {
             item.address: item.string
             for item in extract_strings_from_raw_bytes(
-                read_section_data(pe_file, data_section),
-                base_address=Rva(cast(int, data_section.virtual_address)),
+                data_section.content,
+                base_address=Rva(data_section.virtual_address),
             )
         }
 
         logger.info(f"Found {len(strings)} string-like objects")
 
         logger.info("Searching for cross references...")
         cross_references = find_relative_cross_references(
-            read_section_data(pe_file, code_section),
-            base_address=Rva(cast(int, code_section.virtual_address)),
+            code_section.content,
+            base_address=Rva(code_section.virtual_address),
             addresses=strings,
         )
 
@@ -60,7 +56,7 @@ def patch(patched_file: str | Path, translation_table: list[tuple[str, str]], en
                 if len(translation) <= len(string):
                     # Shorter strings are padded with spaces
                     encoded_translation = translation.encode(encoding).ljust(len(string) + 1) + b"\0"
-                    pe_file.seek(data_section.rva_to_offset(rva))
+                    pe_file.seek(pe.rva_to_offset(rva))
                     pe_file.write(encoded_translation)
                 else:
                     # TODO: implement this case

dfint64_patch/strings_context/extract_strings_with_subs.py

@@ -4,36 +4,35 @@
 
 from collections import defaultdict
 from dataclasses import dataclass
+from io import BufferedReader
 from operator import itemgetter
 from pathlib import Path
-from typing import BinaryIO, NamedTuple
+from typing import NamedTuple
 
+import lief
 from omegaconf import DictConfig
-from peclasses.portable_executable import PortableExecutable
 
-from dfint64_patch.binio import read_section_data
 from dfint64_patch.config import with_config
 from dfint64_patch.cross_references.cross_references_relative import find_relative_cross_references
 from dfint64_patch.extract_strings.from_raw_bytes import ExtractedStringInfo, extract_strings_from_raw_bytes
 from dfint64_patch.extract_subroutines.from_raw_bytes import SubroutineInfo, extract_subroutines, which_subroutine
 from dfint64_patch.type_aliases import Rva
 
 
-def extract_strings_with_xrefs(pe_file: BinaryIO, pe: PortableExecutable) -> dict[ExtractedStringInfo, list[Rva]]:
-    sections = pe.section_table
-    code_section = sections[0]
-    string_section = sections[1]
+def extract_strings_with_xrefs(pe: lief.PE.Binary) -> dict[ExtractedStringInfo, list[Rva]]:
+    code_section = pe.sections[0]
+    string_section = pe.sections[1]
 
     strings = list(
         extract_strings_from_raw_bytes(
-            read_section_data(pe_file, string_section),
-            base_address=string_section.virtual_address,
+            string_section.content,
+            base_address=Rva(string_section.virtual_address),
         ),
     )
 
     cross_references = find_relative_cross_references(
-        read_section_data(pe_file, code_section),
-        base_address=code_section.virtual_address,
+        code_section.content,
+        base_address=Rva(code_section.virtual_address),
         addresses=map(itemgetter(0), strings),
     )
 
@@ -49,18 +48,18 @@ class StringCrossReference(NamedTuple):
     cross_reference: Rva
 
 
-def extract_strings_grouped_by_subs(pe_file: BinaryIO) -> dict[Rva, list[StringCrossReference]]:
-    pe = PortableExecutable(pe_file)
-    sections = pe.section_table
-    code_section = sections[0]
+def extract_strings_grouped_by_subs(pe_file: BufferedReader) -> dict[Rva, list[StringCrossReference]]:
+    pe = lief.PE.parse(pe_file)
+    assert pe is not None
+    code_section = pe.sections[0]
 
-    image_base = pe.optional_header.image_base
+    image_base = pe.optional_header.imagebase
 
-    strings_with_xrefs = extract_strings_with_xrefs(pe_file, pe)
+    strings_with_xrefs = extract_strings_with_xrefs(pe)
 
     subroutines = list(
         extract_subroutines(
-            read_section_data(pe_file, code_section),
+            code_section.content,
             base_offset=code_section.virtual_address,
         )
     )