## ###
# 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 argparse
import platform
import os
import sys
import subprocess
from subprocess import DEVNULL
import sysconfig
from pathlib import Path
from itertools import chain
from typing import List, Dict, Tuple, Optional  

def get_application_properties(install_dir: Path) -> Dict[str, str]:
    app_properties_path: Path = install_dir / 'Ghidra' / 'application.properties'
    props: Dict[str, str] = {}
    with open(app_properties_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#') or line.startswith('!'):
                continue
            key, value = line.split('=', 1)
            if key:
                props[key] = value
    return props

def get_launch_properties(install_dir: Path, dev: bool) -> List[str]:
    if dev:
        launch_properties_path: Path = install_dir / 'Ghidra' / 'RuntimeScripts' / 'support' / 'launch.properties'
    else:
        launch_properties_path: Path = install_dir / 'support' / 'launch.properties'
    props: List[str] = []
    with open(launch_properties_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#') or line.startswith('!'):
                continue
            props.append(line)
    return props

def get_user_settings_dir(install_dir: Path, dev: bool) -> Path:
    app_props: Dict[str, str] = get_application_properties(install_dir)
    app_name: str = app_props['application.name'].replace(' ', '').lower()
    app_version: str = app_props['application.version']
    app_release_name: str = app_props['application.release.name']
    versioned_name: str = f'{app_name}_{app_version}_{app_release_name}'
    if dev:
        ghidra_repos_config = install_dir / 'ghidra.repos.config'
        dir_name = install_dir.parent.name if ghidra_repos_config.is_file() else install_dir.name
        versioned_name += f'_location_{dir_name}'

    # Check for application.settingsdir in launch.properties
    for launch_prop in get_launch_properties(install_dir, dev):
        if launch_prop.startswith('VMARGS=-Dapplication.settingsdir='):
            application_settingsdir = launch_prop[launch_prop.rindex('=')+1:]
            if application_settingsdir:
                return Path(application_settingsdir) / app_name / versioned_name
    
    # Check for XDG_CONFIG_HOME environment variable
    xdg_config_home: str = os.environ.get('XDG_CONFIG_HOME')
    if xdg_config_home:
        return Path(xdg_config_home) / app_name / versioned_name
    
    # Default to platform-specific locations
    if platform.system() == 'Windows':
        return Path(os.environ['APPDATA']) / app_name / versioned_name
    if platform.system() == 'Darwin':
        return Path.home() / 'Library' / app_name / versioned_name
    return Path.home() / '.config' / app_name / versioned_name

def find_supported_python_exe(install_dir: Path, dev: bool) -> List[str]:
    python_cmds = []
    saved_python_cmd = get_saved_python_cmd(install_dir, dev)
    if saved_python_cmd is not None:
        python_cmds.append(saved_python_cmd)
        print("Last used Python executable: " + str(saved_python_cmd))
    
    props: Dict[str, str] = get_application_properties(install_dir)
    prop: str = 'application.python.supported'
    supported: List[str] = [s.strip() for s in props.get(prop, '').split(',')]
    if '' in supported:
        raise ValueError(f'Invalid "{prop}" value in application.properties file')

    python_cmds += list(chain.from_iterable([[f'python{s}'], ['py', f'-{s}']] for s in supported))
    python_cmds += [['python3'], ['python'], ['py']]

    for cmd in python_cmds:
        try:
            result = subprocess.run(cmd + ['-c', 'import sys; print("{0}.{1}".format(*sys.version_info))'], capture_output=True, text=True)
            version = result.stdout.strip()
            if result.returncode == 0 and version in supported:
                return cmd
        except FileNotFoundError:
            pass
        except NotADirectoryError:
            pass
        
    return None
    
def in_venv() -> bool:
    return sys.prefix != sys.base_prefix

def is_externally_managed() -> bool:
    get_default_scheme = 'get_default_scheme'
    if hasattr(sysconfig, get_default_scheme):
        # Python 3.10 and later
        default_scheme = getattr(sysconfig, get_default_scheme)
    else:
        # Python 3.9
        default_scheme = getattr(sysconfig, f'_{get_default_scheme}')
    marker: Path = Path(sysconfig.get_path("stdlib", default_scheme())) / 'EXTERNALLY-MANAGED'
    return marker.is_file()

def get_venv_exe(venv_dir: Path) -> List[str]:
    win_python_cmd: str = str(venv_dir / 'Scripts' / 'python.exe')
    linux_python_cmd: str = str(venv_dir / 'bin' / 'python3')
    return [win_python_cmd] if platform.system() == 'Windows' else [linux_python_cmd]

def get_ghidra_venv(install_dir: Path, dev: bool) -> Path:
    return (install_dir / 'build' if dev else get_user_settings_dir(install_dir, dev)) / 'venv'

def check_module(python_cmd: List[str], module_name: str) -> bool:
    return subprocess.run(python_cmd + ['-c', f'import {module_name}'], stdout=DEVNULL, stderr=DEVNULL).returncode == 0

def check_venv(python_cmd: List[str]) -> bool:
    if check_module(python_cmd, 'ensurepip'):
        return True
    print('A virtual environment cannot be created because "ensurepip" is not available.')
    print('You likely need to install the appropriate "venv" package for your platform and python version.')
    return False

def check_pip(python_cmd: List[str]) -> bool:
    if check_module(python_cmd, 'pip'):
        return True
    print('pip is not present in the current environment.')
    print("Please install pip, remove your virtual environment if applicable, and relaunch PyGhidra.")
    return False

def create_ghidra_venv(python_cmd: List[str], venv_dir: Path) -> None:
    print(f'Creating Ghidra virtual environment at {venv_dir}...')
    subprocess.run(python_cmd + ['-m', 'venv', venv_dir.absolute()])

def version_tuple(v: str) -> Tuple[str, ...]:
    filled = []
    for point in v.split("."):
        filled.append(point.zfill(8))
    return tuple(filled)

def get_package_version(python_cmd: List[str], package: str) -> Optional[str]:
    source = f'import importlib.metadata as m; print(m.version("{package}"))'
    result = subprocess.run(python_cmd + ['-c', source], capture_output=True, text=True)
    return result.stdout.strip() if result.returncode == 0 else None

def get_saved_python_cmd(install_dir: Path, dev: bool) -> List[str]:
    user_settings_dir: Path = get_user_settings_dir(install_dir, dev)
    save_file: Path = user_settings_dir / 'python_command.save'
    if not save_file.is_file():
        return None
    ret = []
    with open(save_file, 'r') as f:
        for line in f:
            ret.append(line.strip())
    return ret

def save_python_cmd(install_dir: Path, python_cmd: List[str], dev: bool) -> None:
    user_settings_dir: Path = get_user_settings_dir(install_dir, dev)
    if not user_settings_dir.is_dir():
        user_settings_dir.mkdir(parents=True, exist_ok=True)
    save_file: Path = user_settings_dir / 'python_command.save'
    with open(save_file, 'w') as f:
        f.write('\n'.join(python_cmd) + '\n')
    
def install(install_dir: Path, python_cmd: List[str], pip_args: List[str], offer_venv: bool) -> List[str]:
    install_choice: str = input('Do you wish to install PyGhidra (y/n)? ')
    if install_choice.lower() in ('y', 'yes'):
        if offer_venv:
            ghidra_venv_choice: str = input('Install into new Ghidra virtual environment (y/n)? ')
            if ghidra_venv_choice.lower() in ('y', 'yes'):
                if not check_venv(python_cmd):
                    return None
                venv_dir = get_ghidra_venv(install_dir, False)
                create_ghidra_venv(python_cmd, venv_dir)
                python_cmd = get_venv_exe(venv_dir)
                print(f'Switching to Ghidra virtual environment: {venv_dir}')
            elif ghidra_venv_choice.lower() in ('n', 'no'):
                system_venv_choice: str = input('Install into system environment (y/n)? ')
                if not system_venv_choice.lower() in ('y', 'yes'):
                    print('Must answer "y" to the prior choices, or launch in an already active virtual environment.')
                    return None
            else:
                print('Please answer yes or no.')
                return None 
        if not check_pip(python_cmd):
            return None
        subprocess.check_call(python_cmd + pip_args)
        return python_cmd
    elif not install_choice.lower() in ('n', 'no'):
        print('Please answer yes or no.')
    return None     

def upgrade(python_cmd: List[str], pip_args: List[str], dist_dir: Path, current_pyghidra_version: str) -> bool:
    included_pyghidra: Path = next(dist_dir.glob('pyghidra-*.whl'), None)
    if included_pyghidra is None:
        print('Warning: included pyghidra wheel was not found', file=sys.stderr)
        return
    included_version = included_pyghidra.name.split('-')[1]
    current_version = current_pyghidra_version
    if version_tuple(included_version) > version_tuple(current_version):
        print(f'PyGhidra upgrade available: {current_version} -> {included_version}')
        if is_externally_managed():
            print(f'Automated upgrade is not supported in an externally managed environment')
            return False
        choice: str = input(f'Do you wish to upgrade PyGhidra {current_version} to {included_version} (y/n)? ')
        if choice.lower() in ('y', 'yes'):
            if not check_pip(python_cmd):
                return False
            pip_args.append('-U')
            subprocess.check_call(python_cmd + pip_args)
            return True
        else:
            print('Skipping upgrade')
            return False
            
def main() -> None:
    # Parse command line arguments
    parser = argparse.ArgumentParser(prog=Path(__file__).name)
    parser.add_argument('install_dir', metavar='<install dir>', help='Ghidra installation directory')
    parser.add_argument('--console', action='store_true', help='Force console launch')
    parser.add_argument('--dev', action='store_true', help='Ghidra development mode')
    parser.add_argument('-H', '--headless', action='store_true', help='Ghidra headless mode')
    args, remaining = parser.parse_known_args()
    
    # Setup variables
    install_dir: Path = Path(os.path.normpath(args.install_dir))
    pyghidra_dir: Path = install_dir / 'Ghidra' / 'Features' / 'PyGhidra'
    dist_dir: Path = pyghidra_dir / 'pypkg' / 'dist'
    venv_dir = get_ghidra_venv(install_dir, args.dev)
    python_cmd: List[str] = find_supported_python_exe(install_dir, args.dev)
      
    if python_cmd is not None:
        print(f'Using Python command: "{" ".join(python_cmd)}"')
    else:
        print('Supported version of Python not found. Check application.properties file.')
        sys.exit(1)
    
    # If headless, force console mode
    if args.headless:
        args.console = True
    
    if args.dev:
        # If in dev mode, launch PyGhidra from the source tree using the development virtual environment
        if not venv_dir.is_dir():
            print('Development virtual environment not found!')
            print('Run "gradle prepPyGhidra" and try again.')
            sys.exit(1)
        python_cmd = get_venv_exe(venv_dir)
        print(f'Switching to Ghidra development virtual environment: {venv_dir}')
    else:
        # If in release mode, offer to install or upgrade PyGhidra before launching from user-controlled environment
        pip_args: List[str] = ['-m', 'pip', 'install', '--no-index', '-f', str(dist_dir), 'pyghidra']

        # Setup the proper execution environment:
        # 1) If we are already in a virtual environment, use that
        # 2) If the Ghidra user settings virtual environment exists, use that
        # 3) If we are "externally managed", automatically create/use the Ghidra user settings virtual environment
        offer_venv = False
        if in_venv():
            # If we are already in a virtual environment, assume that's where the user wants to be
            python_cmd = get_venv_exe(Path(sys.prefix))
            print(f'Using active virtual environment: {sys.prefix}')
        elif os.path.isdir(venv_dir):
            # If the Ghidra user settings venv exists, use that
            python_cmd = get_venv_exe(venv_dir)
            print(f'Switching to Ghidra virtual environment: {venv_dir}')
        elif is_externally_managed():
            print('Externally managed environment detected')
            current_pyghidra_version = get_package_version(python_cmd, 'pyghidra')
            if current_pyghidra_version is None:
                if not check_venv(python_cmd):
                    sys.exit(1)
                create_ghidra_venv(python_cmd, venv_dir)
                python_cmd = get_venv_exe(venv_dir)
                print(f'Switching to Ghidra virtual environment: {venv_dir}')
            else:
                print(f'Using externally managed PyGhidra {current_pyghidra_version}')

        else:
            offer_venv = True

        # If PyGhidra is not installed in the execution environment, offer to install it
        # If it's already installed, offer to upgrade (if applicable)
        current_pyghidra_version = get_package_version(python_cmd, 'pyghidra')
        if current_pyghidra_version is None:
            python_cmd = install(install_dir, python_cmd, pip_args, offer_venv)
            if not python_cmd:
                sys.exit(1)
        else:
            upgrade(python_cmd, pip_args, dist_dir, current_pyghidra_version)

    # Launch PyGhidra
    save_python_cmd(install_dir, python_cmd, args.dev)
    py_args: List[str] = python_cmd + ['-m']
    if args.headless:
        py_args += ['pyghidra.ghidra_launch', '--install-dir', str(install_dir), 'ghidra.app.util.headless.AnalyzeHeadless']
    else:
        py_args += ['pyghidra', '-g', '--install-dir', str(install_dir)]
    if args.console:
        subprocess.call(py_args + remaining)
    else:
        creation_flags = getattr(subprocess, 'CREATE_NO_WINDOW', 0)
        subprocess.Popen(py_args + remaining, creationflags=creation_flags, stdout=DEVNULL, stderr=DEVNULL)

if __name__ == "__main__":
    main()