Refactoring

dfint · Nov 5, 2024 · af4d7d8 · af4d7d8
1 parent 3507dda
commit af4d7d8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/dfint64_patch/extract_strings/cli.py b/dfint64_patch/extract_strings/cli.py
@@ -1,5 +1,4 @@
 import operator
-from collections.abc import Iterable
 from dataclasses import dataclass
 from pathlib import Path
 from typing import BinaryIO, cast
@@ -21,7 +20,7 @@
 from dfint64_patch.utils import maybe_open
 
 
-def extract_strings(pe_file: BinaryIO) -> Iterable[ExtractedStringInfo]:
+def extract_strings(pe_file: BinaryIO) -> list[ExtractedStringInfo]:
     pe = PortableExecutable(pe_file)
 
     sections = pe.section_table

diff --git a/dfint64_patch/strings_context/extract_strings_with_subs.py b/dfint64_patch/strings_context/extract_strings_with_subs.py
@@ -1,5 +1,5 @@
 """
-TODO: Extract strings grouped by subroutines
+Extract strings grouped by subroutines
 """
 from collections import defaultdict
 from dataclasses import dataclass
@@ -48,7 +48,7 @@ class StringCrossReference(NamedTuple):
     cross_reference: Rva
 
 
-def extract_strings_with_subs(pe_file: BinaryIO) -> dict[Rva, list[StringCrossReference]]:
+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]
@@ -87,7 +87,7 @@ class ExtractConfig(DictConfig):
 @with_config(ExtractConfig, ".extract.yaml")
 def main(conf: ExtractConfig) -> None:
     with Path(conf.file_name).open("rb") as pe_file:
-        for subroutine, strings in extract_strings_with_subs(pe_file).items():
+        for subroutine, strings in extract_strings_grouped_by_subs(pe_file).items():
             print(f"sub_{subroutine:x}:")
             for string in strings:
                 print(f"\t{string.string}")