Crusader_Decomp/tools/pyghidra_crusader/cli.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

from .common import (
    DEFAULT_INSTALL_DIR,
    DEFAULT_PROJECT_DIR,
    DEFAULT_PROJECT_NAME,
    DEFAULT_PROGRAM_NAME,
    DEFAULT_FOLDER_PATH,
    ProjectConfig,
    create_function,
    get_function,
    list_root_files,
    open_program,
    open_project,
    remove_function,
    rename_function,
    save_program,
    set_comment,
    transaction,
)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="PyGhidra helpers for the Crusader project."
    )
    parser.add_argument(
        "--install-dir",
        default=str(DEFAULT_INSTALL_DIR),
        help="Ghidra install directory.",
    )
    parser.add_argument(
        "--project-dir",
        default=str(DEFAULT_PROJECT_DIR),
        help="Directory containing the Ghidra project.",
    )
    parser.add_argument(
        "--project-name",
        default=DEFAULT_PROJECT_NAME,
        help="Ghidra project name.",
    )
    parser.add_argument(
        "--program-name",
        default=DEFAULT_PROGRAM_NAME,
        help="Program name inside the project.",
    )
    parser.add_argument(
        "--folder-path",
        default=DEFAULT_FOLDER_PATH,
        help="Project folder path containing the program.",
    )
    parser.add_argument(
        "--restore-project",
        action="store_true",
        help="Restore project tool state while opening the project.",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    subparsers.add_parser(
        "project-files",
        help="List root-level files in the Ghidra project.",
    )

    create_parser = subparsers.add_parser(
        "create-function",
        help="Create a function at an address with an optional explicit body range.",
    )
    create_parser.add_argument("--entry", required=True, help="Function entry address.")
    create_parser.add_argument("--name", required=True, help="New function name.")
    create_parser.add_argument("--body-start", help="Function body start address.")
    create_parser.add_argument("--body-end", help="Function body end address.")
    create_parser.add_argument(
        "--plate-comment",
        help="Optional plate comment to set at the entry address after creation.",
    )

    delete_parser = subparsers.add_parser(
        "delete-function",
        help="Delete a function at an address.",
    )
    delete_parser.add_argument("--entry", required=True, help="Function entry address.")

    rename_parser = subparsers.add_parser(
        "rename-function",
        help="Rename an existing function by entry address.",
    )
    rename_parser.add_argument("--entry", required=True, help="Function entry address.")
    rename_parser.add_argument("--name", required=True, help="New function name.")

    comment_parser = subparsers.add_parser(
        "set-comment",
        help="Set a code-unit comment by address.",
    )
    comment_parser.add_argument("--address", required=True, help="Comment target address.")
    comment_parser.add_argument("--text", required=True, help="Comment text.")
    comment_parser.add_argument(
        "--type",
        choices=["pre", "plate", "eol", "repeatable", "post"],
        default="plate",
        help="Comment type.",
    )

    plan_parser = subparsers.add_parser(
        "apply-plan",
        help="Apply a JSON edit plan containing function and comment operations.",
    )
    plan_parser.add_argument("--plan", required=True, help="Path to the JSON plan file.")
    plan_parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate and print the plan without modifying the project.",
    )

    return parser


def build_config(args: argparse.Namespace) -> ProjectConfig:
    return ProjectConfig(
        install_dir=Path(args.install_dir),
        project_dir=Path(args.project_dir),
        project_name=args.project_name,
        program_name=args.program_name,
        folder_path=args.folder_path,
        restore_project=args.restore_project,
    )
def command_project_files(config: ProjectConfig, _args: argparse.Namespace) -> int:
    project = open_project(config)
    try:
        for name in list_root_files(project):
            print(name)
    finally:
        project.close()
    return 0


def command_create_function(config: ProjectConfig, args: argparse.Namespace) -> int:
    with open_program(config, read_only=False) as (project, program):
        with transaction(program, f"Create function {args.entry}"):
            function = create_function(program, args.entry, args.name, args.body_start, args.body_end)
            if args.plate_comment:
                set_comment(program, args.entry, args.plate_comment, "plate")
        save_program(project, program)
    print(f"created {function.getName()} at {args.entry}")
    return 0


def command_delete_function(config: ProjectConfig, args: argparse.Namespace) -> int:
    with open_program(config, read_only=False) as (project, program):
        with transaction(program, f"Delete function {args.entry}"):
            removed = remove_function(program, args.entry)
        if not removed:
            raise RuntimeError(f"no function removed at {args.entry}")
        save_program(project, program)
    print(f"deleted function at {args.entry}")
    return 0


def command_rename_function(config: ProjectConfig, args: argparse.Namespace) -> int:
    with open_program(config, read_only=False) as (project, program):
        with transaction(program, f"Rename function {args.entry}"):
            function = rename_function(program, args.entry, args.name)
        save_program(project, program)
    print(f"renamed {args.entry} to {function.getName()}")
    return 0


def command_set_comment(config: ProjectConfig, args: argparse.Namespace) -> int:
    with open_program(config, read_only=False) as (project, program):
        with transaction(program, f"Set comment {args.address}"):
            set_comment(program, args.address, args.text, args.type)
        save_program(project, program)
    print(f"set {args.type} comment at {args.address}")
    return 0


def _load_plan(plan_path: str) -> dict:
    with open(plan_path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def _print_plan(plan: dict) -> None:
    print(json.dumps(plan, indent=2, sort_keys=True))


def command_apply_plan(config: ProjectConfig, args: argparse.Namespace) -> int:
    plan = _load_plan(args.plan)
    if args.dry_run:
        _print_plan(plan)
        return 0

    transaction_name = plan.get("transaction", f"Apply plan {args.plan}")
    with open_program(config, read_only=False) as (project, program):
        with transaction(program, transaction_name):
            for entry in plan.get("remove_functions", []):
                removed = remove_function(program, entry)
                if not removed:
                    raise RuntimeError(f"no function removed at {entry}")

            for entry in plan.get("rename_functions", []):
                rename_function(program, entry["entry"], entry["name"])

            for entry in plan.get("create_functions", []):
                create_function(
                    program,
                    entry["entry"],
                    entry["name"],
                    entry.get("body_start"),
                    entry.get("body_end"),
                )
                if entry.get("comment"):
                    set_comment(
                        program,
                        entry["entry"],
                        entry["comment"],
                        entry.get("comment_type", "plate"),
                    )

            for entry in plan.get("comments", []):
                set_comment(
                    program,
                    entry["address"],
                    entry["text"],
                    entry.get("type", "plate"),
                )

            for entry in plan.get("assert_functions", []):
                if get_function(program, entry) is None:
                    raise RuntimeError(f"expected function missing at {entry}")

        save_program(project, program)

    print(f"applied plan {args.plan}")
    return 0


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    config = build_config(args)

    command_map = {
        "project-files": command_project_files,
        "create-function": command_create_function,
        "delete-function": command_delete_function,
        "rename-function": command_rename_function,
        "set-comment": command_set_comment,
        "apply-plan": command_apply_plan,
    }
    return command_map[args.command](config, args)


if __name__ == "__main__":
    raise SystemExit(main())