1# Copyright 2013 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5from telemetry.internal.util import atexit_with_log
6import collections
7import contextlib
8import ctypes
9import logging
10import os
11import platform
12import re
13import socket
14import struct
15import subprocess
16import sys
17import time
18import zipfile
19
20from py_utils import cloud_storage  # pylint: disable=import-error
21
22from telemetry.core import exceptions
23from telemetry.core import os_version as os_version_module
24from telemetry import decorators
25from telemetry.internal.platform import desktop_platform_backend
26from telemetry.internal.platform.power_monitor import msr_power_monitor
27from telemetry.internal.util import path
28
29try:
30  import pywintypes  # pylint: disable=import-error
31  import win32api  # pylint: disable=import-error
32  from win32com.shell import shell  # pylint: disable=no-name-in-module
33  from win32com.shell import shellcon  # pylint: disable=no-name-in-module
34  import win32con  # pylint: disable=import-error
35  import win32file  # pylint: disable=import-error
36  import win32gui  # pylint: disable=import-error
37  import win32pipe  # pylint: disable=import-error
38  import win32process  # pylint: disable=import-error
39  try:
40    import winreg  # pylint: disable=import-error
41  except ImportError:
42    import _winreg as winreg  # pylint: disable=import-error
43  import win32security  # pylint: disable=import-error
44except ImportError:
45  pywintypes = None
46  shell = None
47  shellcon = None
48  win32api = None
49  win32con = None
50  win32file = None
51  win32gui = None
52  win32pipe = None
53  win32process = None
54  win32security = None
55  winreg = None
56
57
58def _InstallWinRing0():
59  """WinRing0 is used for reading MSRs."""
60  executable_dir = os.path.dirname(sys.executable)
61
62  python_is_64_bit = sys.maxsize > 2 ** 32
63  dll_file_name = 'WinRing0x64.dll' if python_is_64_bit else 'WinRing0.dll'
64  dll_path = os.path.join(executable_dir, dll_file_name)
65
66  os_is_64_bit = platform.machine().endswith('64')
67  driver_file_name = 'WinRing0x64.sys' if os_is_64_bit else 'WinRing0.sys'
68  driver_path = os.path.join(executable_dir, driver_file_name)
69
70  # Check for WinRing0 and download if needed.
71  if not (os.path.exists(dll_path) and os.path.exists(driver_path)):
72    win_binary_dir = os.path.join(
73        path.GetTelemetryDir(), 'bin', 'win', 'AMD64')
74    zip_path = os.path.join(win_binary_dir, 'winring0.zip')
75    cloud_storage.GetIfChanged(zip_path, bucket=cloud_storage.PUBLIC_BUCKET)
76    try:
77      with zipfile.ZipFile(zip_path, 'r') as zip_file:
78        error_message = (
79            'Failed to extract %s into %s. If python claims that '
80            'the zip file is locked, this may be a lie. The problem may be '
81            'that python does not have write permissions to the destination '
82            'directory.'
83        )
84        # Install DLL.
85        if not os.path.exists(dll_path):
86          try:
87            zip_file.extract(dll_file_name, executable_dir)
88          except:
89            logging.error(error_message % (dll_file_name, executable_dir))
90            raise
91
92        # Install kernel driver.
93        if not os.path.exists(driver_path):
94          try:
95            zip_file.extract(driver_file_name, executable_dir)
96          except:
97            logging.error(error_message % (driver_file_name, executable_dir))
98            raise
99    finally:
100      os.remove(zip_path)
101
102
103def TerminateProcess(process_handle):
104  if not process_handle:
105    return
106  if win32process.GetExitCodeProcess(process_handle) == win32con.STILL_ACTIVE:
107    win32process.TerminateProcess(process_handle, 0)
108  process_handle.close()
109
110
111class WinPlatformBackend(desktop_platform_backend.DesktopPlatformBackend):
112  def __init__(self):
113    super(WinPlatformBackend, self).__init__()
114    self._msr_server_handle = None
115    self._msr_server_port = None
116    self._power_monitor = msr_power_monitor.MsrPowerMonitorWin(self)
117
118  @classmethod
119  def IsPlatformBackendForHost(cls):
120    return sys.platform == 'win32'
121
122  def __del__(self):
123    self.close()
124
125  def close(self):
126    self.CloseMsrServer()
127
128  def CloseMsrServer(self):
129    if not self._msr_server_handle:
130      return
131
132    TerminateProcess(self._msr_server_handle)
133    self._msr_server_handle = None
134    self._msr_server_port = None
135
136  def IsThermallyThrottled(self):
137    raise NotImplementedError()
138
139  def HasBeenThermallyThrottled(self):
140    raise NotImplementedError()
141
142  def GetSystemCommitCharge(self):
143    performance_info = self._GetPerformanceInfo()
144    return performance_info.CommitTotal * performance_info.PageSize / 1024
145
146  @decorators.Cache
147  def GetSystemTotalPhysicalMemory(self):
148    performance_info = self._GetPerformanceInfo()
149    return performance_info.PhysicalTotal * performance_info.PageSize / 1024
150
151  def GetCpuStats(self, pid):
152    cpu_info = self._GetWin32ProcessInfo(win32process.GetProcessTimes, pid)
153    # Convert 100 nanosecond units to seconds
154    cpu_time = (cpu_info['UserTime'] / 1e7 +
155                cpu_info['KernelTime'] / 1e7)
156    return {'CpuProcessTime': cpu_time}
157
158  def GetCpuTimestamp(self):
159    """Return current timestamp in seconds."""
160    return {'TotalTime': time.time()}
161
162  @decorators.Deprecated(
163      2017, 11, 4,
164      'Clients should use tracing and memory-infra in new Telemetry '
165      'benchmarks. See for context: https://crbug.com/632021')
166  def GetMemoryStats(self, pid):
167    memory_info = self._GetWin32ProcessInfo(
168        win32process.GetProcessMemoryInfo, pid)
169    return {'VM': memory_info['PagefileUsage'],
170            'VMPeak': memory_info['PeakPagefileUsage'],
171            'WorkingSetSize': memory_info['WorkingSetSize'],
172            'WorkingSetSizePeak': memory_info['PeakWorkingSetSize']}
173
174  def KillProcess(self, pid, kill_process_tree=False):
175    # os.kill for Windows is Python 2.7.
176    cmd = ['taskkill', '/F', '/PID', str(pid)]
177    if kill_process_tree:
178      cmd.append('/T')
179    subprocess.Popen(cmd, stdout=subprocess.PIPE,
180                     stderr=subprocess.STDOUT).communicate()
181
182  def GetSystemProcessInfo(self):
183    # [3:] To skip 2 blank lines and header.
184    lines = subprocess.Popen(
185        ['wmic', 'process', 'get',
186         'CommandLine,CreationDate,Name,ParentProcessId,ProcessId',
187         '/format:csv'],
188        stdout=subprocess.PIPE).communicate()[0].splitlines()[3:]
189    process_info = []
190    for line in lines:
191      if not line:
192        continue
193      parts = line.split(',')
194      pi = {}
195      pi['ProcessId'] = int(parts[-1])
196      pi['ParentProcessId'] = int(parts[-2])
197      pi['Name'] = parts[-3]
198      creation_date = None
199      if parts[-4]:
200        creation_date = float(re.split('[+-]', parts[-4])[0])
201      pi['CreationDate'] = creation_date
202      pi['CommandLine'] = ','.join(parts[1:-4])
203      process_info.append(pi)
204    return process_info
205
206  def GetChildPids(self, pid):
207    """Retunds a list of child pids of |pid|."""
208    ppid_map = collections.defaultdict(list)
209    creation_map = {}
210    for pi in self.GetSystemProcessInfo():
211      ppid_map[pi['ParentProcessId']].append(pi['ProcessId'])
212      if pi['CreationDate']:
213        creation_map[pi['ProcessId']] = pi['CreationDate']
214
215    def _InnerGetChildPids(pid):
216      if not pid or pid not in ppid_map:
217        return []
218      ret = [p for p in ppid_map[pid] if creation_map[p] >= creation_map[pid]]
219      for child in ret:
220        if child == pid:
221          continue
222        ret.extend(_InnerGetChildPids(child))
223      return ret
224
225    return _InnerGetChildPids(pid)
226
227  def GetCommandLine(self, pid):
228    for pi in self.GetSystemProcessInfo():
229      if pid == pi['ProcessId']:
230        return pi['CommandLine']
231    raise exceptions.ProcessGoneException()
232
233  @decorators.Cache
234  def GetArchName(self):
235    return platform.machine()
236
237  def GetOSName(self):
238    return 'win'
239
240  @decorators.Cache
241  def GetOSVersionName(self):
242    os_version = platform.uname()[3]
243
244    if os_version.startswith('5.1.'):
245      return os_version_module.XP
246    if os_version.startswith('6.0.'):
247      return os_version_module.VISTA
248    if os_version.startswith('6.1.'):
249      return os_version_module.WIN7
250    # The version of python.exe we commonly use (2.7) is only manifested as
251    # being compatible with Windows versions up to 8. Therefore Windows *lies*
252    # to python about the version number to keep it runnable on Windows 10.
253    key_name = r'Software\Microsoft\Windows NT\CurrentVersion'
254    key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, key_name)
255    try:
256      value, _ = winreg.QueryValueEx(key, 'CurrentMajorVersionNumber')
257    except OSError:
258      value = None
259    finally:
260      key.Close()
261    if value == 10:
262      return os_version_module.WIN10
263    elif os_version.startswith('6.2.'):
264      return os_version_module.WIN8
265    elif os_version.startswith('6.3.'):
266      return os_version_module.WIN81
267    raise NotImplementedError(
268        'Unknown win version: %s, CurrentMajorVersionNumber: %s' %
269        (os_version, value))
270
271  def CanFlushIndividualFilesFromSystemCache(self):
272    return True
273
274  def _GetWin32ProcessInfo(self, func, pid):
275    mask = (win32con.PROCESS_QUERY_INFORMATION |
276            win32con.PROCESS_VM_READ)
277    handle = None
278    try:
279      handle = win32api.OpenProcess(mask, False, pid)
280      return func(handle)
281    except pywintypes.error, e:
282      errcode = e[0]
283      if errcode == 87:
284        raise exceptions.ProcessGoneException()
285      raise
286    finally:
287      if handle:
288        win32api.CloseHandle(handle)
289
290  def _GetPerformanceInfo(self):
291    class PerformanceInfo(ctypes.Structure):
292      """Struct for GetPerformanceInfo() call
293      http://msdn.microsoft.com/en-us/library/ms683210
294      """
295      _fields_ = [('size', ctypes.c_ulong),
296                  ('CommitTotal', ctypes.c_size_t),
297                  ('CommitLimit', ctypes.c_size_t),
298                  ('CommitPeak', ctypes.c_size_t),
299                  ('PhysicalTotal', ctypes.c_size_t),
300                  ('PhysicalAvailable', ctypes.c_size_t),
301                  ('SystemCache', ctypes.c_size_t),
302                  ('KernelTotal', ctypes.c_size_t),
303                  ('KernelPaged', ctypes.c_size_t),
304                  ('KernelNonpaged', ctypes.c_size_t),
305                  ('PageSize', ctypes.c_size_t),
306                  ('HandleCount', ctypes.c_ulong),
307                  ('ProcessCount', ctypes.c_ulong),
308                  ('ThreadCount', ctypes.c_ulong)]
309
310      def __init__(self):
311        self.size = ctypes.sizeof(self)
312        # pylint: disable=bad-super-call
313        super(PerformanceInfo, self).__init__()
314
315    performance_info = PerformanceInfo()
316    ctypes.windll.psapi.GetPerformanceInfo(
317        ctypes.byref(performance_info), performance_info.size)
318    return performance_info
319
320  def IsCurrentProcessElevated(self):
321    if self.GetOSVersionName() < os_version_module.VISTA:
322      # TOKEN_QUERY is not defined before Vista. All processes are elevated.
323      return True
324
325    handle = win32process.GetCurrentProcess()
326    with contextlib.closing(
327        win32security.OpenProcessToken(handle, win32con.TOKEN_QUERY)) as token:
328      return bool(win32security.GetTokenInformation(
329          token, win32security.TokenElevation))
330
331  def LaunchApplication(
332      self, application, parameters=None, elevate_privilege=False):
333    """Launch an application. Returns a PyHANDLE object."""
334
335    parameters = ' '.join(parameters) if parameters else ''
336    if elevate_privilege and not self.IsCurrentProcessElevated():
337      # Use ShellExecuteEx() instead of subprocess.Popen()/CreateProcess() to
338      # elevate privileges. A new console will be created if the new process has
339      # different permissions than this process.
340      proc_info = shell.ShellExecuteEx(
341          fMask=shellcon.SEE_MASK_NOCLOSEPROCESS | shellcon.SEE_MASK_NO_CONSOLE,
342          lpVerb='runas' if elevate_privilege else '',
343          lpFile=application,
344          lpParameters=parameters,
345          nShow=win32con.SW_HIDE)
346      if proc_info['hInstApp'] <= 32:
347        raise Exception('Unable to launch %s' % application)
348      return proc_info['hProcess']
349    else:
350      handle, _, _, _ = win32process.CreateProcess(
351          None, application + ' ' + parameters, None, None, False,
352          win32process.CREATE_NO_WINDOW, None, None, win32process.STARTUPINFO())
353      return handle
354
355  def CanMonitorPower(self):
356    return self._power_monitor.CanMonitorPower()
357
358  def CanMeasurePerApplicationPower(self):
359    return self._power_monitor.CanMeasurePerApplicationPower()
360
361  def StartMonitoringPower(self, browser):
362    self._power_monitor.StartMonitoringPower(browser)
363
364  def StopMonitoringPower(self):
365    return self._power_monitor.StopMonitoringPower()
366
367  def _StartMsrServerIfNeeded(self):
368    if self._msr_server_handle:
369      return
370
371    _InstallWinRing0()
372
373    pipe_name = r"\\.\pipe\msr_server_pipe_{}".format(os.getpid())
374    # Try to open a named pipe to receive a msr port number from server process.
375    pipe = win32pipe.CreateNamedPipe(
376        pipe_name,
377        win32pipe.PIPE_ACCESS_INBOUND,
378        win32pipe.PIPE_TYPE_MESSAGE | win32pipe.PIPE_WAIT,
379        1, 32, 32, 300, None)
380    parameters = (
381        os.path.join(os.path.dirname(__file__), 'msr_server_win.py'),
382        pipe_name,
383    )
384    self._msr_server_handle = self.LaunchApplication(
385        sys.executable, parameters, elevate_privilege=True)
386    if pipe != win32file.INVALID_HANDLE_VALUE:
387      if win32pipe.ConnectNamedPipe(pipe, None) == 0:
388        self._msr_server_port = int(win32file.ReadFile(pipe, 32)[1])
389      win32api.CloseHandle(pipe)
390    # Wait for server to start.
391    try:
392      socket.create_connection(('127.0.0.1', self._msr_server_port), 5).close()
393    except socket.error:
394      self.CloseMsrServer()
395    atexit_with_log.Register(TerminateProcess, self._msr_server_handle)
396
397  def ReadMsr(self, msr_number, start=0, length=64):
398    self._StartMsrServerIfNeeded()
399    if not self._msr_server_handle:
400      raise OSError('Unable to start MSR server.')
401
402    sock = socket.create_connection(('127.0.0.1', self._msr_server_port), 5)
403    try:
404      sock.sendall(struct.pack('I', msr_number))
405      response = sock.recv(8)
406    finally:
407      sock.close()
408    return struct.unpack('Q', response)[0] >> start & ((1 << length) - 1)
409
410  def IsCooperativeShutdownSupported(self):
411    return True
412
413  def CooperativelyShutdown(self, proc, app_name):
414    pid = proc.pid
415
416    # http://timgolden.me.uk/python/win32_how_do_i/
417    #   find-the-window-for-my-subprocess.html
418    #
419    # It seems that intermittently this code manages to find windows
420    # that don't belong to Chrome -- for example, the cmd.exe window
421    # running slave.bat on the tryservers. Try to be careful about
422    # finding only Chrome's windows. This works for both the browser
423    # and content_shell.
424    #
425    # It seems safest to send the WM_CLOSE messages after discovering
426    # all of the sub-process's windows.
427    def find_chrome_windows(hwnd, hwnds):
428      _, win_pid = win32process.GetWindowThreadProcessId(hwnd)
429      if (pid == win_pid and
430          win32gui.IsWindowVisible(hwnd) and
431          win32gui.IsWindowEnabled(hwnd) and
432          win32gui.GetClassName(hwnd).lower().startswith(app_name)):
433        hwnds.append(hwnd)
434      return True
435    hwnds = []
436    win32gui.EnumWindows(find_chrome_windows, hwnds)
437    if hwnds:
438      for hwnd in hwnds:
439        win32gui.SendMessage(hwnd, win32con.WM_CLOSE, 0, 0)
440      return True
441    else:
442      logging.info('Did not find any windows owned by target process')
443    return False
444