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