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 logging
6import os
7import signal
8import subprocess
9import sys
10import tempfile
11
12from telemetry.internal.platform import profiler
13from telemetry.internal.platform.profiler import android_prebuilt_profiler_helper
14
15_TCP_DUMP_BASE_OPTS = ['-i', 'any', '-p', '-s', '0', '-w']
16
17
18class _TCPDumpProfilerAndroid(object):
19  """An internal class to collect TCP dumps on android.
20
21  This profiler uses pre-built binaries from AOSP.
22  See more details in prebuilt/android/README.txt.
23  """
24
25  _DEVICE_DUMP_FILE = '/sdcard/tcpdump_profiles/capture.pcap'
26
27  def __init__(self, device, output_path):
28    self._device = device
29    self._output_path = output_path
30    self._device.RunShellCommand('mkdir -p ' +
31                                 os.path.dirname(self._DEVICE_DUMP_FILE))
32    self._proc = subprocess.Popen(
33        ['adb', '-s', self._device.adb.GetDeviceSerial(),
34         'shell', android_prebuilt_profiler_helper.GetDevicePath('tcpdump')] +
35         _TCP_DUMP_BASE_OPTS +
36         [self._DEVICE_DUMP_FILE])
37
38  def CollectProfile(self):
39    tcpdump_pid = self._device.GetPids('tcpdump')
40    if not tcpdump_pid or not 'tcpdump' in tcpdump_pid:
41      raise Exception('Unable to find TCPDump. Check your device is rooted '
42          'and tcpdump is installed at ' +
43          android_prebuilt_profiler_helper.GetDevicePath('tcpdump'))
44    if len(tcpdump_pid['tcpdump']) > 1:
45      raise Exception(
46          'At most one instance of process tcpdump expected but found pids: '
47          '%s' % tcpdump_pid)
48    tcpdump_pid = int(tcpdump_pid['tcpdump'][0])
49    self._device.RunShellCommand('kill -term ' + tcpdump_pid)
50    self._proc.terminate()
51    host_dump = os.path.join(self._output_path,
52                             os.path.basename(self._DEVICE_DUMP_FILE))
53    try:
54      self._device.PullFile(self._DEVICE_DUMP_FILE, host_dump)
55    except:
56      logging.exception('New exception caused by DeviceUtils conversion')
57      raise
58    print 'TCP dump available at: %s ' % host_dump
59    print 'Use Wireshark to open it.'
60    return host_dump
61
62
63class _TCPDumpProfilerLinux(object):
64  """An internal class to collect TCP dumps on linux desktop."""
65
66  _DUMP_FILE = 'capture.pcap'
67
68  def __init__(self, output_path):
69    if not os.path.exists(output_path):
70      os.makedirs(output_path)
71    self._dump_file = os.path.join(output_path, self._DUMP_FILE)
72    self._tmp_output_file = tempfile.NamedTemporaryFile('w', 0)
73    try:
74      self._proc = subprocess.Popen(
75          ['tcpdump'] + _TCP_DUMP_BASE_OPTS + [self._dump_file],
76          stdout=self._tmp_output_file, stderr=subprocess.STDOUT)
77    except OSError as e:
78      raise Exception('Unable to execute TCPDump, please check your '
79          'installation. ' + str(e))
80
81  def CollectProfile(self):
82    self._proc.send_signal(signal.SIGINT)
83    exit_code = self._proc.wait()
84    try:
85      if exit_code:
86        raise Exception(
87            'tcpdump failed with exit code %d. Output:\n%s' %
88            (exit_code, self._GetStdOut()))
89    finally:
90      self._tmp_output_file.close()
91    print 'TCP dump available at: ', self._dump_file
92    print 'Use Wireshark to open it.'
93    return self._dump_file
94
95  def _GetStdOut(self):
96    self._tmp_output_file.flush()
97    try:
98      with open(self._tmp_output_file.name) as f:
99        return f.read()
100    except IOError:
101      return ''
102
103
104class TCPDumpProfiler(profiler.Profiler):
105  """A Factory to instantiate the platform-specific profiler."""
106  def __init__(self, browser_backend, platform_backend, output_path, state):
107    super(TCPDumpProfiler, self).__init__(
108        browser_backend, platform_backend, output_path, state)
109    if platform_backend.GetOSName() == 'android':
110      android_prebuilt_profiler_helper.InstallOnDevice(
111          browser_backend.device, 'tcpdump')
112      self._platform_profiler = _TCPDumpProfilerAndroid(
113          browser_backend.device, output_path)
114    else:
115      self._platform_profiler = _TCPDumpProfilerLinux(output_path)
116
117  @classmethod
118  def name(cls):
119    return 'tcpdump'
120
121  @classmethod
122  def is_supported(cls, browser_type):
123    if browser_type.startswith('cros'):
124      return False
125    if sys.platform.startswith('linux'):
126      return True
127    return browser_type.startswith('android')
128
129  def CollectProfile(self):
130    return self._platform_profiler.CollectProfile()
131