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 os
10import platform
11import re
12import socket
13import struct
14import subprocess
15import sys
16import time
17import zipfile
18
19from telemetry import decorators
20from telemetry.core import exceptions
21from telemetry.core import util
22from telemetry.core.platform import desktop_platform_backend
23from telemetry.core.platform import platform_backend
24from telemetry.core.platform.power_monitor import msr_power_monitor
25from telemetry.util import cloud_storage
26from telemetry.util import path
27
28try:
29  import pywintypes  # pylint: disable=F0401
30  import win32api  # pylint: disable=F0401
31  from win32com.shell import shell  # pylint: disable=F0401
32  from win32com.shell import shellcon  # pylint: disable=F0401
33  import win32con  # pylint: disable=F0401
34  import win32process  # pylint: disable=F0401
35  import win32security  # pylint: disable=F0401
36except ImportError:
37  pywintypes = None
38  shell = None
39  shellcon = None
40  win32api = None
41  win32con = None
42  win32process = None
43  win32security = None
44
45
46def _InstallWinRing0():
47  """WinRing0 is used for reading MSRs."""
48  executable_dir = os.path.dirname(sys.executable)
49
50  python_is_64_bit = sys.maxsize > 2 ** 32
51  dll_file_name = 'WinRing0x64.dll' if python_is_64_bit else 'WinRing0.dll'
52  dll_path = os.path.join(executable_dir, dll_file_name)
53
54  os_is_64_bit = platform.machine().endswith('64')
55  driver_file_name = 'WinRing0x64.sys' if os_is_64_bit else 'WinRing0.sys'
56  driver_path = os.path.join(executable_dir, driver_file_name)
57
58  # Check for WinRing0 and download if needed.
59  if not (os.path.exists(dll_path) and os.path.exists(driver_path)):
60    win_binary_dir = os.path.join(path.GetTelemetryDir(), 'bin', 'win')
61    zip_path = os.path.join(win_binary_dir, 'winring0.zip')
62    cloud_storage.GetIfChanged(zip_path, bucket=cloud_storage.PUBLIC_BUCKET)
63    try:
64      with zipfile.ZipFile(zip_path, 'r') as zip_file:
65        # Install DLL.
66        if not os.path.exists(dll_path):
67          zip_file.extract(dll_file_name, executable_dir)
68        # Install kernel driver.
69        if not os.path.exists(driver_path):
70          zip_file.extract(driver_file_name, executable_dir)
71    finally:
72      os.remove(zip_path)
73
74
75def TerminateProcess(process_handle):
76  if not process_handle:
77    return
78  if win32process.GetExitCodeProcess(process_handle) == win32con.STILL_ACTIVE:
79    win32process.TerminateProcess(process_handle, 0)
80  process_handle.close()
81
82
83class WinPlatformBackend(desktop_platform_backend.DesktopPlatformBackend):
84  def __init__(self):
85    super(WinPlatformBackend, self).__init__()
86    self._msr_server_handle = None
87    self._msr_server_port = None
88    self._power_monitor = msr_power_monitor.MsrPowerMonitor(self)
89
90  def __del__(self):
91    self.close()
92
93  def close(self):
94    self.CloseMsrServer()
95
96  def CloseMsrServer(self):
97    if not self._msr_server_handle:
98      return
99
100    TerminateProcess(self._msr_server_handle)
101    self._msr_server_handle = None
102    self._msr_server_port = None
103
104  # pylint: disable=W0613
105  def StartRawDisplayFrameRateMeasurement(self):
106    raise NotImplementedError()
107
108  def StopRawDisplayFrameRateMeasurement(self):
109    raise NotImplementedError()
110
111  def GetRawDisplayFrameRateMeasurements(self):
112    raise NotImplementedError()
113
114  def IsThermallyThrottled(self):
115    raise NotImplementedError()
116
117  def HasBeenThermallyThrottled(self):
118    raise NotImplementedError()
119
120  def GetSystemCommitCharge(self):
121    performance_info = self._GetPerformanceInfo()
122    return performance_info.CommitTotal * performance_info.PageSize / 1024
123
124  @decorators.Cache
125  def GetSystemTotalPhysicalMemory(self):
126    performance_info = self._GetPerformanceInfo()
127    return performance_info.PhysicalTotal * performance_info.PageSize / 1024
128
129  def GetCpuStats(self, pid):
130    cpu_info = self._GetWin32ProcessInfo(win32process.GetProcessTimes, pid)
131    # Convert 100 nanosecond units to seconds
132    cpu_time = (cpu_info['UserTime'] / 1e7 +
133                cpu_info['KernelTime'] / 1e7)
134    return {'CpuProcessTime': cpu_time}
135
136  def GetCpuTimestamp(self):
137    """Return current timestamp in seconds."""
138    return {'TotalTime': time.time()}
139
140  def GetMemoryStats(self, pid):
141    memory_info = self._GetWin32ProcessInfo(
142        win32process.GetProcessMemoryInfo, pid)
143    return {'VM': memory_info['PagefileUsage'],
144            'VMPeak': memory_info['PeakPagefileUsage'],
145            'WorkingSetSize': memory_info['WorkingSetSize'],
146            'WorkingSetSizePeak': memory_info['PeakWorkingSetSize']}
147
148  def GetIOStats(self, pid):
149    io_stats = self._GetWin32ProcessInfo(win32process.GetProcessIoCounters, pid)
150    return {'ReadOperationCount': io_stats['ReadOperationCount'],
151            'WriteOperationCount': io_stats['WriteOperationCount'],
152            'ReadTransferCount': io_stats['ReadTransferCount'],
153            'WriteTransferCount': io_stats['WriteTransferCount']}
154
155  def KillProcess(self, pid, kill_process_tree=False):
156    # os.kill for Windows is Python 2.7.
157    cmd = ['taskkill', '/F', '/PID', str(pid)]
158    if kill_process_tree:
159      cmd.append('/T')
160    subprocess.Popen(cmd, stdout=subprocess.PIPE,
161                     stderr=subprocess.STDOUT).communicate()
162
163  def GetSystemProcessInfo(self):
164    # [3:] To skip 2 blank lines and header.
165    lines = subprocess.Popen(
166        ['wmic', 'process', 'get',
167         'CommandLine,CreationDate,Name,ParentProcessId,ProcessId',
168         '/format:csv'],
169        stdout=subprocess.PIPE).communicate()[0].splitlines()[3:]
170    process_info = []
171    for line in lines:
172      if not line:
173        continue
174      parts = line.split(',')
175      pi = {}
176      pi['ProcessId'] = int(parts[-1])
177      pi['ParentProcessId'] = int(parts[-2])
178      pi['Name'] = parts[-3]
179      creation_date = None
180      if parts[-4]:
181        creation_date = float(re.split('[+-]', parts[-4])[0])
182      pi['CreationDate'] = creation_date
183      pi['CommandLine'] = ','.join(parts[1:-4])
184      process_info.append(pi)
185    return process_info
186
187  def GetChildPids(self, pid):
188    """Retunds a list of child pids of |pid|."""
189    ppid_map = collections.defaultdict(list)
190    creation_map = {}
191    for pi in self.GetSystemProcessInfo():
192      ppid_map[pi['ParentProcessId']].append(pi['ProcessId'])
193      if pi['CreationDate']:
194        creation_map[pi['ProcessId']] = pi['CreationDate']
195
196    def _InnerGetChildPids(pid):
197      if not pid or pid not in ppid_map:
198        return []
199      ret = [p for p in ppid_map[pid] if creation_map[p] >= creation_map[pid]]
200      for child in ret:
201        if child == pid:
202          continue
203        ret.extend(_InnerGetChildPids(child))
204      return ret
205
206    return _InnerGetChildPids(pid)
207
208  def GetCommandLine(self, pid):
209    for pi in self.GetSystemProcessInfo():
210      if pid == pi['ProcessId']:
211        return pi['CommandLine']
212    raise exceptions.ProcessGoneException()
213
214  def GetOSName(self):
215    return 'win'
216
217  @decorators.Cache
218  def GetOSVersionName(self):
219    os_version = platform.uname()[3]
220
221    if os_version.startswith('5.1.'):
222      return platform_backend.XP
223    if os_version.startswith('6.0.'):
224      return platform_backend.VISTA
225    if os_version.startswith('6.1.'):
226      return platform_backend.WIN7
227    if os_version.startswith('6.2.'):
228      return platform_backend.WIN8
229
230    raise NotImplementedError('Unknown win version %s.' % os_version)
231
232  def CanFlushIndividualFilesFromSystemCache(self):
233    return True
234
235  def _GetWin32ProcessInfo(self, func, pid):
236    mask = (win32con.PROCESS_QUERY_INFORMATION |
237            win32con.PROCESS_VM_READ)
238    handle = None
239    try:
240      handle = win32api.OpenProcess(mask, False, pid)
241      return func(handle)
242    except pywintypes.error, e:
243      errcode = e[0]
244      if errcode == 87:
245        raise exceptions.ProcessGoneException()
246      raise
247    finally:
248      if handle:
249        win32api.CloseHandle(handle)
250
251  def _GetPerformanceInfo(self):
252    class PerformanceInfo(ctypes.Structure):
253      """Struct for GetPerformanceInfo() call
254      http://msdn.microsoft.com/en-us/library/ms683210
255      """
256      _fields_ = [('size', ctypes.c_ulong),
257                  ('CommitTotal', ctypes.c_size_t),
258                  ('CommitLimit', ctypes.c_size_t),
259                  ('CommitPeak', ctypes.c_size_t),
260                  ('PhysicalTotal', ctypes.c_size_t),
261                  ('PhysicalAvailable', ctypes.c_size_t),
262                  ('SystemCache', ctypes.c_size_t),
263                  ('KernelTotal', ctypes.c_size_t),
264                  ('KernelPaged', ctypes.c_size_t),
265                  ('KernelNonpaged', ctypes.c_size_t),
266                  ('PageSize', ctypes.c_size_t),
267                  ('HandleCount', ctypes.c_ulong),
268                  ('ProcessCount', ctypes.c_ulong),
269                  ('ThreadCount', ctypes.c_ulong)]
270
271      def __init__(self):
272        self.size = ctypes.sizeof(self)
273        super(PerformanceInfo, self).__init__()
274
275    performance_info = PerformanceInfo()
276    ctypes.windll.psapi.GetPerformanceInfo(
277        ctypes.byref(performance_info), performance_info.size)
278    return performance_info
279
280  def IsCurrentProcessElevated(self):
281    if self.GetOSVersionName() < platform_backend.VISTA:
282      # TOKEN_QUERY is not defined before Vista. All processes are elevated.
283      return True
284
285    handle = win32process.GetCurrentProcess()
286    with contextlib.closing(
287        win32security.OpenProcessToken(handle, win32con.TOKEN_QUERY)) as token:
288      return bool(win32security.GetTokenInformation(
289          token, win32security.TokenElevation))
290
291  def LaunchApplication(
292      self, application, parameters=None, elevate_privilege=False):
293    """Launch an application. Returns a PyHANDLE object."""
294
295    parameters = ' '.join(parameters) if parameters else ''
296    if elevate_privilege and not self.IsCurrentProcessElevated():
297      # Use ShellExecuteEx() instead of subprocess.Popen()/CreateProcess() to
298      # elevate privileges. A new console will be created if the new process has
299      # different permissions than this process.
300      proc_info = shell.ShellExecuteEx(
301          fMask=shellcon.SEE_MASK_NOCLOSEPROCESS | shellcon.SEE_MASK_NO_CONSOLE,
302          lpVerb='runas' if elevate_privilege else '',
303          lpFile=application,
304          lpParameters=parameters,
305          nShow=win32con.SW_HIDE)
306      if proc_info['hInstApp'] <= 32:
307        raise Exception('Unable to launch %s' % application)
308      return proc_info['hProcess']
309    else:
310      handle, _, _, _ = win32process.CreateProcess(
311          None, application + ' ' + parameters, None, None, False,
312          win32process.CREATE_NO_WINDOW, None, None, win32process.STARTUPINFO())
313      return handle
314
315  def CanMonitorPower(self):
316    return self._power_monitor.CanMonitorPower()
317
318  def CanMeasurePerApplicationPower(self):
319    return self._power_monitor.CanMeasurePerApplicationPower()
320
321  def StartMonitoringPower(self, browser):
322    self._power_monitor.StartMonitoringPower(browser)
323
324  def StopMonitoringPower(self):
325    return self._power_monitor.StopMonitoringPower()
326
327  def _StartMsrServerIfNeeded(self):
328    if self._msr_server_handle:
329      return
330
331    _InstallWinRing0()
332    self._msr_server_port = util.GetUnreservedAvailableLocalPort()
333    # It might be flaky to get a port number without reserving it atomically,
334    # but if the server process chooses a port, we have no way of getting it.
335    # The stdout of the elevated process isn't accessible.
336    parameters = (
337        os.path.join(os.path.dirname(__file__), 'msr_server_win.py'),
338        str(self._msr_server_port),
339    )
340    self._msr_server_handle = self.LaunchApplication(
341        sys.executable, parameters, elevate_privilege=True)
342    # Wait for server to start.
343    try:
344      socket.create_connection(('127.0.0.1', self._msr_server_port), 5).close()
345    except socket.error:
346      self.CloseMsrServer()
347    atexit.register(TerminateProcess, self._msr_server_handle)
348
349  def ReadMsr(self, msr_number, start=0, length=64):
350    self._StartMsrServerIfNeeded()
351    if not self._msr_server_handle:
352      raise OSError('Unable to start MSR server.')
353
354    sock = socket.create_connection(('127.0.0.1', self._msr_server_port), 0.1)
355    try:
356      sock.sendall(struct.pack('I', msr_number))
357      response = sock.recv(8)
358    finally:
359      sock.close()
360    return struct.unpack('Q', response)[0] >> start & ((1 << length) - 1)
361