1424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)# Copyright 2013 The Chromium Authors. All rights reserved.
2424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)# Use of this source code is governed by a BSD-style license that can be
3424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)# found in the LICENSE file.
4424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
5424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)"""Common helper module for working with Chrome's processes and windows."""
6424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
758537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)import ctypes
858537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)import pywintypes
9424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)import re
1058537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)import win32con
11424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)import win32gui
12424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)import win32process
13424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
1458537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)
1558537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)def GetProcessIDAndPathPairs():
1658537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)  """Returns a list of 2-tuples of (process id, process path).
1758537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)
1858537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)  This is needed because psutil is not available on Windows slave machines (see:
1958537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)  http://crbug.com/257696).
2058537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)  TODO(sukolsak): Use psutil.process_iter() once it becomes available.
2158537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)  """
2258537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)  process_id_and_path_pairs = []
2358537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)  for process_id in win32process.EnumProcesses():
2458537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)    process_handle = ctypes.windll.kernel32.OpenProcess(
2558537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)        win32con.PROCESS_QUERY_INFORMATION | win32con.PROCESS_VM_READ, False,
2658537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)        process_id)
2758537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)    if not process_handle:
2858537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)      continue
2958537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)    try:
3058537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)      process_path = win32process.GetModuleFileNameEx(process_handle, 0)
3158537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)      process_id_and_path_pairs.append((process_id, process_path))
3258537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)    except pywintypes.error:
3358537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)      # It's normal that some processes are not accessible.
3458537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)      pass
3558537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)  return process_id_and_path_pairs
36424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
37424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
38424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)def GetProcessIDs(process_path):
39424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  """Returns a list of IDs of processes whose path is |process_path|.
40424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
41424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  Args:
42424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    process_path: The path to the process.
43424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
44424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  Returns:
45424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    A list of process IDs.
46424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  """
4758537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)  return [pid for (pid, path) in GetProcessIDAndPathPairs() if
4858537e28ecd584eab876aee8be7156509866d23aTorne (Richard Coles)          path == process_path]
49424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
50424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
51424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)def GetWindowHandles(process_ids):
52424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  """Returns a list of handles of windows owned by processes in |process_ids|.
53424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
54424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  Args:
55424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    process_ids: A list of process IDs.
56424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
57424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  Returns:
58424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    A list of handles of windows owned by processes in |process_ids|.
59424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  """
60424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  hwnds = []
61424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  def EnumerateWindowCallback(hwnd, _):
62424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    _, found_process_id = win32process.GetWindowThreadProcessId(hwnd)
63424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    if found_process_id in process_ids and win32gui.IsWindowVisible(hwnd):
64424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)      hwnds.append(hwnd)
65424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  # Enumerate all the top-level windows and call the callback with the hwnd as
66424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  # the first parameter.
67424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  win32gui.EnumWindows(EnumerateWindowCallback, None)
68424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  return hwnds
69424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
70424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
71424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)def WindowExists(process_ids, class_pattern):
72424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  """Returns whether there exists a window with the specified criteria.
73424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
74424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  This method returns whether there exists a window that is owned by a process
75424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  in |process_ids| and has a class name that matches |class_pattern|.
76424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
77424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  Args:
78424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    process_ids: A list of process IDs.
79424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    class_pattern: The regular expression pattern of the window class name.
80424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)
81424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  Returns:
82424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    A boolean indicating whether such window exists.
83424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  """
84424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  for hwnd in GetWindowHandles(process_ids):
85424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)    if re.match(class_pattern, win32gui.GetClassName(hwnd)):
86424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)      return True
87424c4d7b64af9d0d8fd9624f381f469654d5e3d2Torne (Richard Coles)  return False
88