diff --git a/tools/ghidra_scripts/import_functions_and_types_from_pdb.py b/tools/ghidra_scripts/import_functions_and_types_from_pdb.py index 54bb7e9f..2ac843b1 100644 --- a/tools/ghidra_scripts/import_functions_and_types_from_pdb.py +++ b/tools/ghidra_scripts/import_functions_and_types_from_pdb.py @@ -21,9 +21,7 @@ # Disable spurious warnings in vscode / pylance # pyright: reportMissingModuleSource=false -from enum import Enum import importlib -from dataclasses import dataclass, field import logging.handlers import sys import logging @@ -53,43 +51,8 @@ def reload_module(module: str): reload_module("lego_util.statistics") -from lego_util.statistics import Statistics - - -@dataclass -class Globals: - verbose: bool - loglevel: int - running_from_ghidra: bool = False - # statistics - statistics: Statistics = field(default_factory=Statistics) - - -class SupportedModules(Enum): - LEGO1 = 1 - BETA10 = 2 - - def orig_filename(self): - if self == self.LEGO1: - return "LEGO1.DLL" - return "BETA10.DLL" - - def recomp_filename_without_extension(self): - # in case we want to support more functions - return "LEGO1" - - def build_dir_name(self): - if self == self.BETA10: - return "build_debug" - return "build" - - -# hard-coded settings that we don't want to prompt in Ghidra every time -GLOBALS = Globals( - verbose=False, - # loglevel=logging.INFO, - loglevel=logging.DEBUG, -) +reload_module("lego_util.globals") +from lego_util.globals import GLOBALS, SupportedModules def setup_logging(): @@ -232,24 +195,24 @@ def main(): origfile_name = getProgramFile().getName() if origfile_name == "LEGO1.DLL": - module = SupportedModules.LEGO1 + GLOBALS.module = SupportedModules.LEGO1 elif origfile_name in ["LEGO1D.DLL", "BETA10.DLL"]: - module = SupportedModules.BETA10 + GLOBALS.module = SupportedModules.BETA10 else: raise Lego1Exception( f"Unsupported file name in import script: {origfile_name}" ) - else: - module = SupportedModules.LEGO1 - logger.info("Importing file: %s", module.orig_filename()) + logger.info("Importing file: %s", GLOBALS.module.orig_filename()) repo_root = get_repository_root() - origfile_path = repo_root.joinpath("legobin").joinpath(module.orig_filename()) - build_directory = repo_root.joinpath(module.build_dir_name()) - recompiledfile_name = f"{module.recomp_filename_without_extension()}.DLL" + origfile_path = repo_root.joinpath("legobin").joinpath( + GLOBALS.module.orig_filename() + ) + build_directory = repo_root.joinpath(GLOBALS.module.build_dir_name()) + recompiledfile_name = f"{GLOBALS.module.recomp_filename_without_extension()}.DLL" recompiledfile_path = build_directory.joinpath(recompiledfile_name) - pdbfile_name = f"{module.recomp_filename_without_extension()}.PDB" + pdbfile_name = f"{GLOBALS.module.recomp_filename_without_extension()}.PDB" pdbfile_path = build_directory.joinpath(pdbfile_name) if not GLOBALS.verbose: diff --git a/tools/ghidra_scripts/lego_util/ghidra_helper.py b/tools/ghidra_scripts/lego_util/ghidra_helper.py index 24e70c65..91717efd 100644 --- a/tools/ghidra_scripts/lego_util/ghidra_helper.py +++ b/tools/ghidra_scripts/lego_util/ghidra_helper.py @@ -8,6 +8,7 @@ from lego_util.exceptions import ( TypeNotFoundInGhidraError, MultipleTypesFoundInGhidraError, ) +from lego_util.globals import GLOBALS, SupportedModules # Disable spurious warnings in vscode / pylance # pyright: reportMissingModuleSource=false @@ -105,6 +106,12 @@ def sanitize_name(name: str) -> str: .replace("`", "'") ) + # Importing function names like `FUN_10001234` into BETA10 can be confusing + # because Ghidra's auto-generated functions look exactly the same. + # Therefore, such function names are replaced by `LEGO_10001234` in the BETA10 import. + if GLOBALS.module == SupportedModules.BETA10: + new_name = re.sub(r"FUN_([0-9a-f]{8})", r"LEGO1_\1", new_name) + if "<" in name: new_name = "_template_" + new_name diff --git a/tools/ghidra_scripts/lego_util/globals.py b/tools/ghidra_scripts/lego_util/globals.py new file mode 100644 index 00000000..541b1106 --- /dev/null +++ b/tools/ghidra_scripts/lego_util/globals.py @@ -0,0 +1,42 @@ +import logging +from enum import Enum +from dataclasses import dataclass, field +from lego_util.statistics import Statistics + + +class SupportedModules(Enum): + LEGO1 = 1 + BETA10 = 2 + + def orig_filename(self): + if self == self.LEGO1: + return "LEGO1.DLL" + return "BETA10.DLL" + + def recomp_filename_without_extension(self): + # in case we want to support more functions + return "LEGO1" + + def build_dir_name(self): + if self == self.BETA10: + return "build_debug" + return "build" + + +@dataclass +class Globals: + verbose: bool + loglevel: int + module: SupportedModules + running_from_ghidra: bool = False + # statistics + statistics: Statistics = field(default_factory=Statistics) + + +# hard-coded settings that we don't want to prompt in Ghidra every time +GLOBALS = Globals( + verbose=False, + # loglevel=logging.INFO, + loglevel=logging.DEBUG, + module=SupportedModules.LEGO1, # this default value will be used when run outside of Ghidra +)