diff --git a/search_offsets/search.py b/search_offsets/search.py index b347e64..00202ec 100644 --- a/search_offsets/search.py +++ b/search_offsets/search.py @@ -36,23 +36,24 @@ def search(path: str, patterns: list[Pattern]) -> Mapping[str, list[int]]: return found -def print_found(section_table: SectionTable, pattern_names: Iterable[str], found: Mapping[str, int]) -> None: +def process_found(section_table: SectionTable, pattern_names: Iterable[Pattern], found: Mapping[str, int]) -> dict[str, int]: """ Print offsets of the found patterns. """ for pattern in pattern_names: - if not found[pattern]: - print(f"{pattern}: NOT FOUND") + pattern_name = pattern.name + if not found[pattern_name]: + yield pattern_name, None continue - for i, offset in enumerate(found[pattern], -1): + for i, offset in enumerate(found[pattern_name], -1): suffix = "" if i < 0 else f"_{i}" - name = pattern + suffix + name = pattern_name + suffix if name == "addchar_0": name = "addchar_top" rva = section_table.offset_to_rva(offset) - print(f"{name} = 0x{rva:X}") + yield name, rva @dataclass @@ -75,10 +76,15 @@ def main(config: DictConfig) -> None: section_table = pe.section_table found = search(config.path, patterns) - print_found(section_table, map(str, patterns), found) + processed = dict(process_found(section_table, patterns, found)) + for key, value in processed.items(): + if value is None: + print(f"{key}: NOT FOUND") + else: + print(f"{key} = 0x{value:X}") if config.get("version_name", None): - result = render_template(found, checksum=pe.file_header.timedate_stamp, version_name=config.version_name) + result = render_template(processed, checksum=pe.file_header.timedate_stamp, version_name=config.version_name) file_name = f"offsets_{config.version_name.replace(' ', '_')}.toml" (root_dir / file_name).write_text(result, encoding="utf-8") diff --git a/search_offsets/templates/offsets.toml.jinja b/search_offsets/templates/offsets.toml.jinja index 6d3cc5c..3856914 100644 --- a/search_offsets/templates/offsets.toml.jinja +++ b/search_offsets/templates/offsets.toml.jinja @@ -4,15 +4,15 @@ version = "{{ version_name }}" checksum = {{ checksum | hex }} [offsets] -string_copy_n = {{ string_copy_n[0] | hex }} -string_append_n = {{ string_append_n[0] | hex }} -std_string_ctor = {{ std_string_ctor[0] | hex }} -addst = {{ addst[0] | hex }} -addst_top = {{ addst_top[0] | hex }} -addst_flag = {{ addst_flag[0] | hex }} -standardstringentry = {{ standardstringentry[0] | hex }} -simplify_string = {{ simplify_string[0] | hex }} -upper_case_string = {{ upper_case_string[0] | hex }} -lower_case_string = {{ lower_case_string[0] | hex }} -capitalize_string_words = {{ capitalize_string_words[0] | hex }} -capitalize_string_first_word = {{ capitalize_string_first_word[0] | hex }} +string_copy_n = {{ string_copy_n | hex }} +string_append_n = {{ string_append_n | hex }} +std_string_ctor = {{ std_string_ctor | hex }} +addst = {{ addst | hex }} +addst_top = {{ addst_top | hex }} +addst_flag = {{ addst_flag | hex }} +standardstringentry = {{ standardstringentry | hex }} +simplify_string = {{ simplify_string | hex }} +upper_case_string = {{ upper_case_string | hex }} +lower_case_string = {{ lower_case_string | hex }} +capitalize_string_words = {{ capitalize_string_words | hex }} +capitalize_string_first_word = {{ capitalize_string_first_word | hex }}