## ###
# IP: GHIDRA
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##
import sys
import contextlib
from typing import Union, TYPE_CHECKING, Tuple, List, Callable, Generator, Any, Optional

from pyghidra.converters import *  # pylint: disable=wildcard-import, unused-wildcard-import

if TYPE_CHECKING:
    from pyghidra.launcher import PyGhidraLauncher
    from ghidra.app.plugin.core.analysis import AutoAnalysisManager
    from ghidra.program.model.listing import Program
    from ghidra.program.util import GhidraProgramUtilities
    from ghidra.framework.model import Project, DomainFile
    from ghidra.framework.options import Options
    from ghidra.formats.gfilesystem import GFileSystem
    from ghidra.formats.gfilesystem import FSUtilities
    from ghidra.util.task import TaskMonitor
    from ghidra.app.script import GhidraScript
    from ghidra.app.util.importer import ProgramLoader
    from ghidra.pyghidra import PyGhidraTaskMonitor
    from generic.jar import ResourceFile
    from java.lang import Object # type:ignore @UnresolvedImport

def start(verbose=False, *, install_dir: Optional[Path] = None) -> "PyGhidraLauncher":
    """
    Starts the JVM and fully initializes Ghidra in Headless mode.

    :param verbose: Enable verbose output during JVM startup (Defaults to False)
    :param install_dir: The path to the Ghidra installation directory.
        (Defaults to the GHIDRA_INSTALL_DIR environment variable or "lastrun" file)
    :return: The PyGhidraLauncher used to start the JVM
    """
    from pyghidra.launcher import HeadlessPyGhidraLauncher
    launcher = HeadlessPyGhidraLauncher(verbose=verbose,  install_dir=install_dir)
    launcher.start()
    return launcher


def started() -> bool:
    """
    Whether the PyGhidraLauncher has already started.
    """
    from pyghidra.launcher import PyGhidraLauncher
    return PyGhidraLauncher.has_launched()

def open_project(
        path: Union[str, Path],
        name: str,
        create: bool = False
) -> "Project":
    """
    Opens the Ghidra project at the given location, optionally creating it if it doesn't exist.

    :param path: Path of Ghidra project parent directory.
    :param name: Name of Ghidra project to open/create.
    :param create: Whether to create the project if it doesn't exist
    :return: A Ghidra "Project" object.
    :raises FileNotFoundError: If the project to open was not found and it shouldn't be created.
    """
    from ghidra.framework.model import ProjectLocator
    from ghidra.pyghidra import PyGhidraProjectManager
    
    projectLocator = ProjectLocator(path, name);
    projectManager = PyGhidraProjectManager()
    if projectLocator.exists():
        return projectManager.openProject(projectLocator, False, True);
    elif create:
        return projectManager.createProject(projectLocator, None, True)
    raise FileNotFoundError(f'Project "{name}" not found at "{path}"!')

def open_filesystem(
        path: Union[str, Path]
    ) -> "GFileSystem":
    """
    Opens a filesystem in Ghidra.

    :param path: Path of filesystem to open in Ghidra.
    :return: A Ghidra "GFileSystem" object.
    :raises ValueError: If the filesystem to open is not supported by Ghidra.
    """
    from java.io import File # type:ignore @UnresolvedImport
    from ghidra.formats.gfilesystem import FileSystemService
    
    service = FileSystemService.getInstance()
    fsrl = service.getLocalFS().getLocalFSRL(File(path))
    fs = service.openFileSystemContainer(fsrl, task_monitor())
    if fs is None:
        raise ValueError(f'"{fsrl}" is not a supported GFileSystem!')
    return fs

def consume_program(
        project: "Project", 
        path: Union[str, Path],
        consumer: Optional[Any] = None
    ) -> Tuple["Program", "Object"]:
    """
    Gets the Ghidra program from the given project with the given project path. The returned program
    must be manually released when it is no longer needed.

    :param project: The Ghidra project that has the program.
    :param path: The project path of the program (should start with "/")
    :param consumer: An optional reference to the Java object "consuming" the returned program, used
        to ensure the underlying DomainObject is only closed when every consumer is done with it. If
        a consumer is not provided, one will be generated by this function.
    :return: A 2-element tuple containing the program and a consumer object that must be used to
        release the program when finished with it (i.e., program.release(consumer). If a consumer
        object was provided, the same consumer object is returned. Otherwise, a new consumer object
        is created and returned.
    :raises FileNotFoundError: If the path does not exist in the project.
    :raises pyghidra.ProgramTypeError: If the path in the project exists but is not a Program.
    """
    from ghidra.program.model.listing import Program
    from java.lang import Object # type:ignore @UnresolvedImport
    if consumer is None:
        consumer = Object()
    project_data = project.getProjectData()
    df = project_data.getFile(path)
    if df is None:
        raise FileNotFoundError(f'"{path}" does not exist in the Project')
    dobj = df.getDomainObject(consumer, True, False, task_monitor())
    program_cls = Program.class_
    if not program_cls.isAssignableFrom(dobj.getClass()):
        dobj.release(consumer)
        raise ProgramTypeError(f'"{path}" exists but is not a Program')
    return dobj, consumer

@contextlib.contextmanager
def program_context(
        project: "Project", 
        path: Union[str, Path],
    ) -> Generator["Program", None, None]:
    """
    Gets the Ghidra program from the given project with the given project path. The returned
    program's resource cleanup is performed by a context manager.

    :param project: The Ghidra project that has the program.
    :param path: The project path of the program (should start with "/").
    :return: The Ghidra program.
    :raises FileNotFoundError: If the path does not exist in the project.
    :raises pyghidra.ProgramTypeError: If the path in the project exists but is not a Program.
    """
    program, consumer = consume_program(project, path)    
    try:
        yield program
    finally:
        program.release(consumer)

def analyze(
        program: "Program", 
        monitor: Optional["TaskMonitor"] = None
    ) -> str:
    """
    Analyzes the given program.

    :param program: The Ghidra program to analyze.
    :return: The analysis log.
    """
    from ghidra.app.script import GhidraScriptUtil
    from ghidra.program.util import GhidraProgramUtilities
    from ghidra.app.plugin.core.analysis import AutoAnalysisManager
    
    if monitor is None:
        monitor = task_monitor()
        
    with transaction(program, "Analyze"):
        GhidraScriptUtil.acquireBundleHostReference()
        try:
            mgr: AutoAnalysisManager = AutoAnalysisManager.getAnalysisManager(program);
            mgr.initializeOptions();
            mgr.reAnalyzeAll(None);
            mgr_log = ""
            def get_log(manager, _is_cancelled):
                nonlocal mgr_log
                mgr_log += manager.getMessageLog().toString()
            mgr.addListener(get_log)
            mgr.startAnalysis(monitor, True); # yields to analysis
            if not monitor.isCancelled():
                monitor.cancel()
            GhidraProgramUtilities.markProgramAnalyzed(program)
            return mgr_log
        finally:
            GhidraScriptUtil.releaseBundleHostReference()
    
def ghidra_script(
        path: Union[str, Path],
        project: "Project",
        program: Optional["Program"] = None,
        script_args: List[str] = [],
        echo_stdout = True,
        echo_stderr = True
    ) -> Tuple[str, str]:
    """
    Runs any type of GhidraScript (Java, PyGhidra, Jython, etc).

    :param path: The GhidraScript's path.
    :param project: The Ghidra project to run the GhidraScript in.
    :param program: An optional Ghidra program that the GhidraScript will see as its "currentProgram".
    :param script_args An optional list of arguments to pass to the GhidraScript.
    :param echo_stdout: Whether or not to echo the GhidraScript's standard output.
    :param echo_stderr: Whether or not to echo the GhidraScript's standard error.
    :return: A 2 element tuple consisting of the GhidraScript's standard output and standard error.
    """
    from generic.jar import ResourceFile
    from ghidra.app.script import GhidraScriptUtil, GhidraState, ScriptControls
    from java.io import File, PrintWriter, StringWriter # type:ignore @UnresolvedImport
    from java.lang import System # type:ignore @UnresolvedImport

    GhidraScriptUtil.acquireBundleHostReference()
    try:
        source_file = ResourceFile(File(path))
        if not source_file.exists():
            raise TypeError(f'"{str(source_file)}" was not found')
        provider = GhidraScriptUtil.getProvider(source_file)
        if provider is None:
            raise TypeError(f'"{path}" is not a supported GhidraScript')
        script = provider.getScriptInstance(source_file,  PrintWriter(System.out))
        if script is None:
            raise TypeError(f'"{str(source_file)}" was not found')
        state = GhidraState(None, project, program, None, None, None)
        stdout_string_writer = StringWriter()
        stderr_string_writer = StringWriter()
        controls = ScriptControls(
            PrintWriter(stdout_string_writer, True),
            PrintWriter(stderr_string_writer, True),
            task_monitor()
        )
        script.setScriptArgs(script_args)
        script.execute(state, controls)
        stdout_str = str(stdout_string_writer)
        stderr_str = str(stderr_string_writer)
        if echo_stdout:
            sys.stdout.write(stdout_str)
            sys.stdout.flush()
        if echo_stderr:
            sys.stderr.write(stderr_str)
            sys.stderr.flush()
        return stdout_str, stderr_str
    finally:
        GhidraScriptUtil.releaseBundleHostReference()

@contextlib.contextmanager
def transaction(
        program: "Program",
        description: str = "Unnamed Transaction"
    ) -> Generator[int, None, None]:
    """
    Creates a context for running a Ghidra transaction.

    :param program: The Ghidra program that will be affected.
    :param description: The transaction description
    :return: The transaction ID.
    """
    transaction_id = program.startTransaction(description)
    success = True
    try:
        yield transaction_id
    except:
        success = False
        raise
    finally:
        program.endTransaction(transaction_id, success)

def analysis_properties(program: "Program") -> "Options":
    """
    Convenience function to get the Ghidra "Program.ANALYSIS_PROPERTIES" options.

    :return: the Ghidra "Program.ANALYSIS_PROPERTIES" options.
    """
    from ghidra.program.model.listing import Program
    from ghidra.app.plugin.core.analysis import AutoAnalysisManager
    AutoAnalysisManager.getAnalysisManager(program).initializeOptions()
    return program.getOptions(Program.ANALYSIS_PROPERTIES)

def program_info(program: "Program") -> "Options":
    """
    Convenience function to get the Ghidra "Program.PROGRAM_INFO" options.

    :return: the Ghidra "Program.PROGRAM_INFO" options.
    """
    from ghidra.program.model.listing import Program
    return program.getOptions(Program.PROGRAM_INFO)

def program_loader() -> "ProgramLoader.Builder":
    """
    Convenience function to get a Ghidra "ProgramLoader.Builder" object.

    :return: A Ghidra "ProgramLoader.Builder" object.
    """
    from ghidra.app.util.importer import ProgramLoader
    return ProgramLoader.builder()

def task_monitor(
        timeout: Optional[int] = None
    ) -> "TaskMonitor":
    """
    Convenience function to get a "TaskMonitor" object.

    :param timeout: An optional number of seconds to wait before canceling the monitor.
    :return: A "TaskMonitor" object.
    """
    from ghidra.util.task import TaskMonitor
    from ghidra.pyghidra import PyGhidraTaskMonitor
    from jpype.types import JInt
    if timeout is None:
        return TaskMonitor.DUMMY
    else:
        return PyGhidraTaskMonitor(JInt(timeout), None)

def walk_project(
        project: "Project",
        callback: Callable[["DomainFile"], None],
        start: Union[str, Path] = "/",
        file_filter: Callable[["DomainFile"], bool] = lambda _f: True
    ):
    """
    Walks the the given Ghidra project, calling the provided function when each domain file is 
    encountered.

    :param project: The Ghidra project to walk.
    :param callback: The callback to process each domain file.
    :param start: An optional starting project folder path.
    :param file_filter: A filter used to limit what domain files get processed.
    :raises FileNotFoundError: If the starting folder is not found in the project.
    """
    from ghidra.framework.model import ProjectDataUtils
    start_folder = project.projectData.getFolder(start)
    if start_folder is None:
        raise FileNotFoundError(f'Starting folder "{start}" does not exist in the Project')
    for file in ProjectDataUtils.descendantFiles(start_folder):
        if file_filter(file):
            callback(file)

def walk_programs(
        project: "Project",
        callback: Callable[["DomainFile", "Program"], None],
        start: Union[str, Path] = "/",
        program_filter: Callable[["DomainFile", "Program"], bool] = lambda _f, _p: True
    ):
    """
    Walks the the given Ghidra project, calling the provided function when each program is 
    encountered. Non-programs in the project are skipped.

    :param project: The Ghidra project to walk.
    :param callback: The callback to process each program.
    :param start: An optional starting project folder path.
    :param program_filter: A filter used to limit what programs get processed.
    :raises FileNotFoundError: If the starting folder is not found in the project.
    """
    def process(file: "DomainFile"):
        try:
            with program_context(project, file.getPathname()) as program:
                if program_filter(file, program):
                    callback(file, program)
        except ProgramTypeError:
            pass # skip over non-programs 
    
    walk_project(project, process, start=start)
    
class ProgramTypeError(TypeError):
    """Custom exception for when a Program was expected but not received."""
    pass