Crusader_Decomp/tools/pyghidra_crusader/common.py

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
import os


REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_INSTALL_DIR = Path(
    os.environ.get("GHIDRA_INSTALL_DIR", r"I:\Apps\ghidra_11.3.2_PUBLIC")
)
DEFAULT_PROJECT_DIR = REPO_ROOT
DEFAULT_PROJECT_NAME = "Crusader"
DEFAULT_PROGRAM_NAME = "CRUSADER-RAW.EXE"
DEFAULT_FOLDER_PATH = "/"


@dataclass(frozen=True)
class ProjectConfig:
    install_dir: Path = DEFAULT_INSTALL_DIR
    project_dir: Path = DEFAULT_PROJECT_DIR
    project_name: str = DEFAULT_PROJECT_NAME
    program_name: str = DEFAULT_PROGRAM_NAME
    folder_path: str = DEFAULT_FOLDER_PATH
    restore_project: bool = False


def ensure_pyghidra_started(install_dir: Path | None = None):
    import pyghidra

    resolved_dir = Path(install_dir or DEFAULT_INSTALL_DIR)
    if not pyghidra.started():
        pyghidra.start(install_dir=resolved_dir)
    return pyghidra


def parse_address_text(address_text: str) -> int:
    text = address_text.strip()
    if ":" in text:
        segment_text, offset_text = text.split(":", 1)
        return (int(segment_text, 16) << 16) + int(offset_text, 16)
    return int(text, 0)


def to_address(program, address_text: str):
    address_space = program.getAddressFactory().getDefaultAddressSpace()
    return address_space.getAddress(parse_address_text(address_text))


def format_project_error(config: ProjectConfig, exc: Exception) -> RuntimeError:
    lock_path = config.project_dir / f"{config.project_name}.lock"
    details = [
        f"unable to open project '{config.project_name}' in '{config.project_dir}'",
        str(exc),
    ]
    if lock_path.exists():
        details.append(
            f"project lock present at '{lock_path}'; close Ghidra or work on a project copy for write operations"
        )
    return RuntimeError("; ".join(details))


def open_project(config: ProjectConfig):
    ensure_pyghidra_started(config.install_dir)

    from ghidra.base.project import GhidraProject

    try:
        return GhidraProject.openProject(
            str(config.project_dir),
            config.project_name,
            config.restore_project,
        )
    except Exception as exc:  # pragma: no cover - depends on local Ghidra state
        raise format_project_error(config, exc) from exc


def _candidate_folder_paths(folder_path: str) -> list[str]:
    candidates = [folder_path]
    for fallback in ("/", "\\", ""):
        if fallback not in candidates:
            candidates.append(fallback)
    return candidates


@contextmanager
def open_program(config: ProjectConfig, read_only: bool):
    project = open_project(config)
    program = None
    last_error = None
    try:
        for folder_path in _candidate_folder_paths(config.folder_path):
            try:
                program = project.openProgram(folder_path, config.program_name, read_only)
                break
            except Exception as exc:  # pragma: no cover - depends on local Ghidra state
                last_error = exc
        if program is None:
            raise RuntimeError(
                f"unable to open program '{config.program_name}' from project '{config.project_name}': {last_error}"
            )
        yield project, program
    finally:
        if project is not None:
            if program is not None:
                project.close(program)
            project.close()


@contextmanager
def transaction(program, description: str):
    transaction_id = program.startTransaction(description)
    commit = False
    try:
        yield
        commit = True
    finally:
        program.endTransaction(transaction_id, commit)


def list_root_files(project) -> list[str]:
    return [domain_file.getName() for domain_file in project.getRootFolder().getFiles()]


def get_function(program, entry_text: str):
    return program.getFunctionManager().getFunctionAt(to_address(program, entry_text))


def create_function(program, entry_text: str, name: str, body_start: str | None, body_end: str | None):
    from ghidra.program.model.address import AddressSet
    from ghidra.program.model.symbol import SourceType

    entry_address = to_address(program, entry_text)
    body_start_address = to_address(program, body_start or entry_text)
    body_end_address = to_address(program, body_end or entry_text)
    body = AddressSet(body_start_address, body_end_address)
    return program.getFunctionManager().createFunction(
        name,
        entry_address,
        body,
        SourceType.USER_DEFINED,
    )


def remove_function(program, entry_text: str) -> bool:
    return bool(program.getFunctionManager().removeFunction(to_address(program, entry_text)))


def rename_function(program, entry_text: str, new_name: str):
    from ghidra.program.model.symbol import SourceType

    function = get_function(program, entry_text)
    if function is None:
        raise ValueError(f"no function found at {entry_text}")
    function.setName(new_name, SourceType.USER_DEFINED)
    return function


def set_comment(program, address_text: str, comment: str, comment_type: str):
    from ghidra.program.model.listing import CodeUnit

    comment_types = {
        "pre": CodeUnit.PRE_COMMENT,
        "plate": CodeUnit.PLATE_COMMENT,
        "eol": CodeUnit.EOL_COMMENT,
        "repeatable": CodeUnit.REPEATABLE_COMMENT,
        "post": CodeUnit.POST_COMMENT,
    }
    if comment_type not in comment_types:
        raise ValueError(f"unsupported comment type: {comment_type}")

    listing = program.getListing()
    code_unit = listing.getCodeUnitAt(to_address(program, address_text))
    if code_unit is None:
        raise ValueError(f"no code unit found at {address_text}")
    code_unit.setComment(comment_types[comment_type], comment)


def save_program(project, program):
    project.save(program)