chromium/chrome/updater/test/service/win/ui.py

# Copyright 2021 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging
import time

import pythoncom
import pywintypes
import win32api
import win32com
import win32com.client
import win32con
import win32gui
import win32process
import winerror


class _MessageQueueAttacher(object):
    """Wrapper class for message queue attachment."""

    def __enter__(self):
        """Attaches the current thread to the foreground window's message queue.

        This is an old and well known exploit used to bypass Windows Focus
        rules:
        http://www.google.com/search?q=attachthreadinput+setforegroundwindow
        """
        self._active_thread_id = 0
        active_hwnd = win32gui.GetForegroundWindow()
        if not active_hwnd:
            logging.warning('No active window is found.')
            return
        current_thread_id = win32api.GetCurrentThreadId()
        active_thread_id, _ = win32process.GetWindowThreadProcessId(
            active_hwnd)
        win32process.AttachThreadInput(current_thread_id, active_thread_id, 1)
        logging.info('Attached current thread input %s to active thread: %s',
                     current_thread_id, active_thread_id)
        self._active_thread_id = active_thread_id

    def __exit__(self, unused_type, unused_value, unused_traceback):
        """Detaches the current thread from the active thread's message queue.
        """
        if not self._active_thread_id:
            return
        current_thread_id = win32api.GetCurrentThreadId()
        win32process.AttachThreadInput(current_thread_id,
                                       self._active_thread_id, 0)
        logging.info('Detached current thread input %s from thread: %s',
                     current_thread_id, self._active_thread_id)


def SetForegroundWindow(hwnd):
    """Brings the given window to foreground.

    Args:
        hwnd: Handle of the window to bring to the foreground.
    """
    with _MessageQueueAttacher():
        return bool(win32gui.SetForegroundWindow(hwnd))


def FindWindowsWithText(parent, text_to_search):
    """Finds windows with given text.

    Args:
        parent: Handle to the parent window whose child windows are to be
                searched.
        text_to_search: Substring to search within Windows text,
                        case-insensitive.

    Returns:
        A list of HWND that match the search condition.
    """

    class WindowFoundHandler(object):
        """Callback class for window enumeration."""

        def __init__(self, text_to_search):
            self.result = []
            self._text_to_search = text_to_search

        def Process(self, handle):
            """Callback function when enumerating a window.

      Args:
          handle: HWND to the enumerated window.
      """
            text = win32gui.GetWindowText(handle).lower()
            text_to_search = self._text_to_search.lower()

            if text_to_search in text:
                self.result.append(handle)

    def WinFoundCallback(hwnd, window_found_handler):
        window_found_handler.Process(hwnd)

    handler = WindowFoundHandler(text_to_search)
    try:
        win32gui.EnumChildWindows(parent, WinFoundCallback, handler)
    except pywintypes.error as e:
        logging.info('Error while searching [%s], error: [%s]', text_to_search,
                     e)

    return handler.result


def FindWindowsWithTitle(title_to_search):
    """Finds windows with given title.

    Args:
        title_to_search: Window title substring to search, case-insensitive.

    Returns:
        A list of HWND that match the search condition.
    """
    desktop_handle = None
    return FindWindowsWithText(desktop_handle, title_to_search)


def FindWindow(title, class_name, parent=0, child_after=0):
    """Finds a window of a given title and class.

    Args:
        title: Title of the window to search.
        class_name: Class name of the window to search.
        parent: Handle to the parent window whose child windows are to be
                searched.
        child_after: HWND to a child window. Search begins with the next child
          window in the Z order.

    Returns:
        Handle of the found window, or 0 if not found.
    """
    hwnd = 0
    try:
        hwnd = win32gui.FindWindowEx(int(parent), int(child_after), class_name,
                                     title)
    except win32gui.error as err:
        if err[0] == winrror.ERROR_INVALID_WINDOW_HANDLE:  # Could be closed.
            pass
        elif err[0] != winrror.ERROR_FILE_NOT_FOUND:
            raise err
    if hwnd:
        win32gui.FlashWindow(hwnd, True)
    return hwnd


def FindWindowWithTitleAndText(title, text):
    """Checks if the any window has given title, and child window has the text.

    Args:
        title: Expected window title.
        text: Expected window text substring(in any child window),
              case-insensitive.

    Returns:
        A list how HWND that meets the search condition.
    """
    hwnds_title_matched = FindWindowsWithTitle(title)
    if not hwnds_title_matched:
        logging.info('No window has title: [%s].', title)

    hwnds = []
    for hwnd in hwnds_title_matched:
        if FindWindowsWithText(hwnd, text):
            hwnds.append(hwnd)
    return hwnds


def WaitForWindow(title, class_name, timeout=30):
    """Waits for window with given title and class to appear.

    Args:
        title: Windows title to search.
        class_name: Class name of the window to search.
        timeout: How long should wait before give up.

    Returns:
        A tuple of (HWND, title) of the found window, or (None, None) otherwise.
    """
    logging.info('ui.WaitForWindow("%s", "%s") for %s seconds', title,
                 class_name, timeout)
    hwnd = None
    start = time.perf_counter()
    stop = start + int(timeout)

    while time.perf_counter() < stop:
        hwnd = FindWindow(title, class_name)
        if hwnd:
            elapsed = time.perf_counter() - start
            logging.info('Window ["%s"] found in %f seconds', title, elapsed)
            return (hwnd, title)
        logging.info('Window with title [%s] has not appeared yet.', title)
        time.sleep(0.5)

    logging.warning('WARNING: (%s,"%s") not found within %f seconds', title,
                    class_name, timeout)
    return (None, None)


def ClickButton(button_hwnd):
    """Clicks a button window by sending a BM_CLICK message.

    Per http://msdn2.microsoft.com/en-us/library/bb775985.aspx
    "If the button is in a dialog box and the dialog box is not active, the
    BM_CLICK message might fail. To ensure success in this situation, call
    the SetActiveWindow function to activate the dialog box before sending
    the BM_CLICK message to the button."

    Args:
        button_hwnd: HWND to the button to be clicked.
    """
    previous_active_window = win32gui.SetActiveWindow(button_hwnd)
    win32gui.PostMessage(button_hwnd, win32con.BM_CLICK, 0, 0)
    if previous_active_window:
        win32gui.SetActiveWindow(previous_active_window)


def ClickChildButtonWithText(parent_hwnd, button_text):
    """Clicks a child button window with given text.

    Args:
        parent_hwnd: HWND of the button parent.
        button_text: Button windows title.

    Returns:
        Whether button is clicked.
    """
    button_hwnd = FindWindow(button_text, 'Button', parent_hwnd)
    if button_hwnd:
        logging.debug('Found child button with: %s', button_text)
        ClickButton(button_hwnd)
        return True

    logging.debug('No button with [%s] found.', button_text)
    return False


def SendKeyToWindow(hwnd, key_to_press):
    """Sends a key press to the window.

    This is a blocking call until all keys are pressed.

    Args:
        hwnd: handle of  the window to press key.
        key_to_press: Actual key to send to window. eg.'{Enter}'.
                      See `window_shell` documentation for key definitions.
    """
    try:
        window_shell = win32com.client.Dispatch('WScript.Shell')
        SetForegroundWindow(hwnd)
        window_shell.AppActivate(str(win32gui.GetForegroundWindow()))
        window_shell.SendKeys(key_to_press)
        logging.info('Sent %s to window %x.', key_to_press, hwnd)
    except pywintypes.error as err:
        logging.error('Failed to press key: %s', err)
        raise
    except pythoncom.com_error as err:
        logging.error('COM exception occurred: %s, is CoInitialize() called?',
                      err)
        raise