o3de/Tools/LyTestTools/ly_test_tools/environment/process_utils.py

"""
Copyright (c) Contributors to the Open 3D Engine Project

SPDX-License-Identifier: Apache-2.0 OR MIT

Process management functions, to supplement normal use of psutil and subprocess
"""
import logging
import os
import psutil
import subprocess
import ctypes

import ly_test_tools.environment.waiter as waiter
from ly_test_tools import WINDOWS, MAC

logger = logging.getLogger(__name__)
_PROCESS_OUTPUT_ENCODING = 'utf-8'

# Default list of processes names to kill
LY_PROCESS_KILL_LIST = [
    'CrySCompileServer', 'Editor',
    'Profiler', 'RemoteConsole',
    'rc'  # Resource Compiler
]


def kill_processes_named(names, ignore_extensions=False):
    """
    Kills all processes with a given name

    :param names: string process name, or list of strings of process name
    :param ignore_extensions: ignore trailing file extension
    """
    if not names:
        return

    names = [names] if isinstance(names, str) else names

    if ignore_extensions:
        names = [_remove_extension(name) for name in names]

    # remove any blank names, which may empty the list
    names = list(filter(lambda x: not x.isspace(), names))
    if not names:
        return

    logger.info(f"Killing all processes named {names}")
    process_list_to_kill = []
    for process in _safe_get_processes(['name', 'pid']):
        try:
            proc_name = process.name()
        except psutil.AccessDenied:
            logger.info(f"Process {process} permissions error during kill_processes_named()", exc_info=True)
            continue
        except psutil.ProcessLookupError:
            logger.debug(f"Process {process} could not be killed during kill_processes_named() and was likely already stopped", exc_info=True)
            continue
        except psutil.NoSuchProcess:
            logger.debug(f"Process '{process}' was active when list of processes was requested but it was not found "
                         f"during kill_processes_named()", exc_info=True)
            continue

        if ignore_extensions:
            proc_name = _remove_extension(proc_name)

        if proc_name in names:
            logger.debug(f"Found process with name {proc_name}.  Attempting to kill...")
            process_list_to_kill.append(process)

    _safe_kill_process_list(process_list_to_kill)


def kill_processes_started_from(path):
    """
    Kills all processes started from a given directory or executable

    :param path: path to application or directory
    """
    logger.info(f"Killing processes started from '{path}'")
    if os.path.exists(path):
        process_list = []
        for process in _safe_get_processes():
            try:
                process_path = process.exe()
            except (psutil.AccessDenied, psutil.NoSuchProcess):
                continue

            if process_path.lower().startswith(path.lower()):
                process_list.append(process)

        _safe_kill_process_list(process_list)
    else:
        logger.warning(f"Path:'{path}' not found")


def kill_processes_with_name_not_started_from(name, path):
    """
    Kills all processes with a given name that NOT started from a directory or executable

    :param name: name of application to look for
    :param path: path where process shouldn't have started from
    """
    path = os.path.join(os.getcwd(), os.path.normpath(path)).lower()
    logger.info(f"Killing processes with name:'{name}' not started from '{path}'")
    if os.path.exists(path):
        proccesses_to_kill = []
        for process in _safe_get_processes(["name", "pid"]):
            try:
                process_path = process.exe()
            except (psutil.AccessDenied, psutil.NoSuchProcess) as ex:
                continue

            process_name = os.path.splitext(os.path.basename(process_path))[0]

            if process_name == os.path.basename(name) and not os.path.dirname(process_path.lower()) == path:
                logger.info("%s -> %s" % (os.path.dirname(process_path.lower()), path))
                proccesses_to_kill.append(process)

        _safe_kill_process_list(proccesses_to_kill)
    else:
        logger.warning(f"Path:'{path}' not found")


def kill_process_with_pid(pid, raise_on_missing=False):
    """
    Kills the process with the specified pid

    :param pid: the pid of the process to kill
    :param raise_on_missing: if set to True, raise RuntimeError if the process does not already exist
    """
    if pid is None:
        logger.warning("Killing process id of 'None' will terminate the current python process!")
    logger.info(f"Killing processes with id '{pid}'")
    process = psutil.Process(pid)
    if process.is_running():
        _safe_kill_process(process)
    elif raise_on_missing:
        message = f"Process with id {pid} was not present"
        logger.error(message)
        raise RuntimeError(message)


def process_exists(name, ignore_extensions=False):
    """
    Determines whether a process with the given name exists

    :param name: process name
    :param ignore_extensions: ignore trailing file extension
    :return: A boolean determining whether the process is alive or not
    """
    name = name.lower()
    if ignore_extensions:
        name = _remove_extension(name)
    if name.isspace():
        return False
    for process in _safe_get_processes(["name"]):
        try:
            proc_name = process.name().lower()
        except psutil.NoSuchProcess as e:
            logger.debug(f"Process '{process}' was active when list of processes was requested but it was not found "
                         f"during process_exists()", exc_info=True)
            continue
        except psutil.AccessDenied as e:
            logger.info(f"Permissions issue on {process} during process_exists check", exc_info=True)
            continue
        if ignore_extensions:
            proc_name = _remove_extension(proc_name)
        if proc_name == name:
            return True
    return False


def process_is_unresponsive(name):
    """
    Check if the specified process is unresponsive.

    Mac warning: this method assumes that a process is not responsive if it is sleeping or waiting, this is true for
    'active' applications, but may not be the case for power optimized applications.

    :param name: the name of the process to check
    :return: True if the specified process is unresponsive and False otherwise
    """
    if WINDOWS:
        output = check_output(['tasklist',
                               '/FI', f'IMAGENAME eq {name}',
                               '/FI', 'STATUS eq NOT RESPONDING'])
        output = output.split(os.linesep)
        for line in output:
            if line and name.startswith(line.split()[0]):
                logger.debug(f"Process '{name}' was unresponsive.")
                logger.debug(line)
                return True
        logger.debug(f"Process '{name}' was not unresponsive.")
        return False
    elif MAC:
        cmd = ["ps", "-axc", "-o", "command,state"]
        output = check_output(cmd)
        for line in output.splitlines()[1:]:
            info = [l.strip() for l in line.split(" ") if l.strip() != '']
            state = info[-1]
            pname = " ".join(info[0:-1])

            if pname == name:
                logger.debug(f"{pname}: {state}")
                if "R" not in state:
                    logger.debug(f"Process {name} was unresponsive.")
                    return True
        logger.debug(f"Process '{name}' was not unresponsive.")
        return False
    else:
        raise NotImplementedError('Only Windows and Mac hosts are supported.')


def check_output(command, **kwargs):
    """
    Forwards arguments to subprocess.check_output so better error messages can be displayed upon failure.

    If you need the stderr output from a failed process then pass in stderr=subprocess.STDOUT as a kwarg.

    :param command: A list of the command to execute and its arguments as split by whitespace.
    :param kwargs: Keyword args forwarded to subprocess.check_output.
    :return: Output from the command if it succeeds.
    """
    cmd_string = command
    if type(command) == list:
        cmd_string = ' '.join(command)

    logger.info(f'Executing "check_output({cmd_string})"')
    try:
        output = subprocess.check_output(command, **kwargs).decode(_PROCESS_OUTPUT_ENCODING)
    except subprocess.CalledProcessError as e:
        logger.error(f'Command "{cmd_string}" failed with returncode {e.returncode}, output:\n{e.output}')
        raise
    logger.info(f'Successfully executed "check_output({cmd_string})"')
    return output


def safe_check_output(command, **kwargs):
    """
    Forwards arguments to subprocess.check_output so better error messages can be displayed upon failure.
    This function eats the subprocess.CalledProcessError exception upon command failure and returns the output.

    If you need the stderr output from a failed process then pass in stderr=subprocess.STDOUT as a kwarg.

    :param command: A list of the command to execute and its arguments as split by whitespace.
    :param kwargs: Keyword args forwarded to subprocess.check_output.
    :return: Output from the command regardless of its return value.
    """
    cmd_string = command
    if type(command) == list:
        cmd_string = ' '.join(command)

    logger.info(f'Executing "check_output({cmd_string})"')
    try:
        output = subprocess.check_output(command, **kwargs).decode(_PROCESS_OUTPUT_ENCODING)
    except subprocess.CalledProcessError as e:
        output = e.output
        logger.warning(f'Command "{cmd_string}" failed with returncode {e.returncode}, output:\n{e.output}')
    else:
        logger.info(f'Successfully executed "check_output({cmd_string})"')
    return output


def check_call(command, **kwargs):
    """
    Forwards arguments to subprocess.check_call so better error messages can be displayed upon failure.

    :param command: A list of the command to execute and its arguments as if split by whitespace.
    :param kwargs: Keyword args forwarded to subprocess.check_call.
    :return: An exitcode of 0 if the call succeeds.
    """
    cmd_string = command
    if type(command) == list:
        cmd_string = ' '.join(command)

    logger.info(f'Executing "check_call({cmd_string})"')
    try:
        subprocess.check_call(command, **kwargs)
    except subprocess.CalledProcessError as e:
        logger.error(f'Command "{cmd_string}" failed with returncode {e.returncode}')
        raise
    logger.info(f'Successfully executed "check_call({cmd_string})"')
    return 0


def safe_check_call(command, **kwargs):
    """
    Forwards arguments to subprocess.check_call so better error messages can be displayed upon failure.
    This function eats the subprocess.CalledProcessError exception upon command failure and returns the exit code.

    :param command: A list of the command to execute and its arguments as if split by whitespace.
    :param kwargs: Keyword args forwarded to subprocess.check_call.
    :return: An exitcode of 0 if the call succeeds, otherwise the exitcode returned from the failed subprocess call.
    """
    cmd_string = command
    if type(command) == list:
        cmd_string = ' '.join(command)

    logger.info(f'Executing "check_call({cmd_string})"')
    try:
        subprocess.check_call(command, **kwargs)
    except subprocess.CalledProcessError as e:
        logger.warning(f'Command "{cmd_string}" failed with returncode {e.returncode}')
        return e.returncode
    else:
        logger.info(f'Successfully executed "check_call({cmd_string})"')
    return 0


def _safe_get_processes(attrs=None):
    """
    Returns the process iterator without raising an error if the process list changes

    :return: The process iterator
    """
    processes = None
    max_attempts = 10
    for _ in range(max_attempts):
        try:
            processes = psutil.process_iter(attrs)
            break
        except (psutil.Error, RuntimeError):
            logger.debug("Unexpected error", exc_info=True)
            continue
    return processes


def _safe_kill_process(proc):
    """
    Kills a given process without raising an error

    :param proc: The process to kill
    """
    try:
        logger.info(f"Terminating process '{proc.name()}' with id '{proc.pid}'")
        _terminate_and_confirm_dead(proc)
    except psutil.AccessDenied:
        logger.warning("Termination failed, Access Denied", exc_info=True)
    except psutil.NoSuchProcess:
        logger.debug("Termination request ignored, process was already terminated during iteration", exc_info=True)
    except Exception:  # purposefully broad
        logger.warning("Unexpected exception while terminating process", exc_info=True)

def _safe_kill_process_list(proc_list):
    """
    Kills a given process without raising an error

    :param proc_list: The process list to kill
    """

    def on_terminate(proc):
        print(f"process '{proc.name()}' with id '{proc.pid}' terminated with exit code {proc.returncode}")

    for proc in proc_list:
        try:
            logger.info(f"Terminating process '{proc.name()}' with id '{proc.pid}'")
            proc.kill()
        except psutil.AccessDenied:
            logger.warning("Termination failed, Access Denied", exc_info=True)
        except psutil.NoSuchProcess:
            logger.debug("Termination request ignored, process was already terminated during iteration", exc_info=True)
        except Exception:  # purposefully broad
            logger.warning("Unexpected exception while terminating process", exc_info=True)

    try:
        psutil.wait_procs(proc_list, timeout=30, callback=on_terminate)
    except Exception:  # purposefully broad
        logger.warning("Unexpected exception while waiting for process to terminate", exc_info=True)


def _terminate_and_confirm_dead(proc):
    """
    Kills a process and waits for the process to stop running.

    :param proc: A process to kill, and wait for proper termination
    """
    def killed():
        return not proc.is_running()

    proc.kill()
    waiter.wait_for(killed, exc=RuntimeError("Process did not terminate after kill command"))


def _remove_extension(filename):
    """
    Returns a file name without its extension

    :param filename: The name of a file
    :return: The name of the file without the extension
    """
    return filename.rsplit(".", 1)[0]


def close_windows_process(pid, timeout=20, raise_on_missing=False):
    # type: (int, int, bool) -> None
    """
    Closes a window using the windows api and checks the return code. An error will be raised if the window hasn't
    closed after the timeout duration.
    Note: This is for Windows only and will fail on any other OS

    :param pid: The process id of the process window to close
    :param timeout: How long to wait for the window to close (seconds)
    :return: None
    :param pid: the pid of the process to kill
    :param raise_on_missing: if set to True, raise RuntimeError if the process does not already exist
    """
    if not WINDOWS:
        raise NotImplementedError("close_windows_process() is only implemented on Windows.")

    if pid is None:
        raise TypeError("Cannot close window with pid of None")

    if not psutil.Process(pid).is_running():
        if raise_on_missing:
            message = f"Process with id {pid} was unexpectedly not present"
            logger.error(message)
            raise RuntimeError(message)
        else:
            logger.warning(f"Process with id {pid} was not present but option raise_on_missing is disabled. Unless "
                           f"a matching process gets opened, calling close_windows_process will spin until its timeout")

    # Gain access to windows api
    user32 = ctypes.windll.user32
    # Set up C data types and function params
    WNDENUMPROC = ctypes.WINFUNCTYPE(ctypes.wintypes.BOOL,
                                     ctypes.wintypes.HWND,
                                     ctypes.wintypes.LPARAM)
    user32.EnumWindows.argtypes = [
        WNDENUMPROC,
        ctypes.wintypes.LPARAM]
    user32.GetWindowTextLengthW.argtypes = [
        ctypes.wintypes.HWND]

    # This is called for each process window
    def _close_matched_process_window(hwnd, _):
        # type: (ctypes.wintypes.HWND, int) -> bool
        """
        EnumWindows() takes a function argument that will return True/False to keep iterating or not.

        Checks the windows handle's pid against the given pid. If they match, then the window will be closed and
        returns False.

        :param hwnd: A windows handle to check against the pid
        :param _: Unused buffer parameter
        :return: False if process was found and closed, else True
        """
        # Get the process id of the handle
        lpdw_process_id = ctypes.c_ulong()
        user32.GetWindowThreadProcessId(hwnd, ctypes.byref(lpdw_process_id))
        process_id = lpdw_process_id.value
        # Compare to the process id
        if pid == process_id:
            # Close the window
            WM_CLOSE = 16  # System message for closing window: 0x10
            user32.PostMessageA(hwnd, WM_CLOSE, 0, 0)
            # Found window for process id, stop iterating
            return False
        # Process not found, keep looping
        return True

    # Call the function on all of the handles
    close_process_func = WNDENUMPROC(_close_matched_process_window)
    user32.EnumWindows(close_process_func, 0)

    # Wait for asyncronous termination
    waiter.wait_for(lambda: pid not in psutil.pids(), timeout=timeout,
                    exc=TimeoutError(f"Process {pid} never terminated"))