tcpdump_profiler.py revision 4e180b6a0b4720a9b8e9e959a882386f690f08ff 
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 os
6import signal
7import subprocess
8import sys
9import tempfile
10
11from telemetry.core.platform import profiler
12from telemetry.core.platform.profiler import android_prebuilt_profiler_helper
13
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, adb, output_path):
28    self._adb = adb
29    self._output_path = output_path
30    self._adb.RunShellCommand('mkdir -p ' +
31                              os.path.dirname(self._DEVICE_DUMP_FILE))
32    self._proc = subprocess.Popen(
33        ['adb', '-s', self._adb.device(),
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._adb.ExtractPid('tcpdump')
40    if not tcpdump_pid or not tcpdump_pid[0]:
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    self._adb.RunShellCommand('kill -term ' + tcpdump_pid[0])
45    self._proc.terminate()
46    host_dump = os.path.join(self._output_path,
47                             os.path.basename(self._DEVICE_DUMP_FILE))
48    self._adb.Adb().Adb().Pull(self._DEVICE_DUMP_FILE, host_dump)
49    print 'TCP dump available at: %s ' % host_dump
50    print 'Use Wireshark to open it.'
51    return host_dump
52
53
54class _TCPDumpProfilerLinux(object):
55  """An internal class to collect TCP dumps on linux desktop."""
56
57  _DUMP_FILE = 'capture.pcap'
58
59  def __init__(self, output_path):
60    if not os.path.exists(output_path):
61      os.makedirs(output_path)
62    self._dump_file = os.path.join(output_path, self._DUMP_FILE)
63    self._tmp_output_file = tempfile.NamedTemporaryFile('w', 0)
64    try:
65      self._proc = subprocess.Popen(
66          ['tcpdump'] + _TCP_DUMP_BASE_OPTS + [self._dump_file],
67          stdout=self._tmp_output_file, stderr=subprocess.STDOUT)
68    except OSError as e:
69      raise Exception('Unable to execute TCPDump, please check your '
70          'installation. ' + str(e))
71
72  def CollectProfile(self):
73    self._proc.send_signal(signal.SIGINT)
74    exit_code = self._proc.wait()
75    try:
76      if exit_code:
77        raise Exception(
78            'tcpdump failed with exit code %d. Output:\n%s' %
79            (exit_code, self._GetStdOut()))
80    finally:
81      self._tmp_output_file.close()
82    print 'TCP dump available at: ', self._dump_file
83    print 'Use Wireshark to open it.'
84    return self._dump_file
85
86  def _GetStdOut(self):
87    self._tmp_output_file.flush()
88    try:
89      with open(self._tmp_output_file.name) as f:
90        return f.read()
91    except IOError:
92      return ''
93
94
95class TCPDumpProfiler(profiler.Profiler):
96  """A Factory to instantiate the platform-specific profiler."""
97  def __init__(self, browser_backend, platform_backend, output_path, state):
98    super(TCPDumpProfiler, self).__init__(
99        browser_backend, platform_backend, output_path, state)
100    if platform_backend.GetOSName() == 'android':
101      android_prebuilt_profiler_helper.InstallOnDevice(
102          browser_backend.adb, 'tcpdump')
103      self._platform_profiler = _TCPDumpProfilerAndroid(
104          browser_backend.adb, output_path)
105    else:
106      self._platform_profiler = _TCPDumpProfilerLinux(output_path)
107
108  @classmethod
109  def name(cls):
110    return 'tcpdump'
111
112  @classmethod
113  def is_supported(cls, browser_type):
114    if browser_type.startswith('cros'):
115      return False
116    if sys.platform.startswith('linux'):
117      return True
118    return browser_type.startswith('android')
119
120  def CollectProfile(self):
121    return self._platform_profiler.CollectProfile()
122