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.chrome import android_browser_finder
12from telemetry.core.platform import profiler
13
14
15_TCP_DUMP_BASE_OPTS = ['-i', 'any', '-p', '-s', '0', '-w']
16
17class _TCPDumpProfilerAndroid(object):
18  """An internal class to collect TCP dumps on android."""
19
20  _TCP_DUMP = '/data/local/tmp/tcpdump'
21  _DEVICE_DUMP_FILE = '/sdcard/capture.pcap'
22
23  def __init__(self, adb, output_path):
24    self._adb = adb
25    self._output_path = output_path
26    self._proc = subprocess.Popen(
27        ['adb', '-s', self._adb.device(),
28         'shell', self._TCP_DUMP] + _TCP_DUMP_BASE_OPTS +
29         [self._DEVICE_DUMP_FILE])
30
31  def CollectProfile(self):
32    tcpdump_pid = self._adb.ExtractPid('tcpdump')
33    if not tcpdump_pid or not tcpdump_pid[0]:
34      raise Exception('Unable to find TCPDump. Check your device is rooted '
35          'and tcpdump is installed at ' + self._TCP_DUMP)
36    self._adb.RunShellCommand('kill -term ' + tcpdump_pid[0])
37    self._proc.terminate()
38    host_dump = os.path.join(self._output_path,
39                             os.path.basename(self._DEVICE_DUMP_FILE))
40    self._adb.Adb().Adb().Pull(self._DEVICE_DUMP_FILE, host_dump)
41    print 'TCP dump available at: %s ' % host_dump
42    print 'Use Wireshark to open it.'
43
44
45class _TCPDumpProfilerLinux(object):
46  """An internal class to collect TCP dumps on linux desktop."""
47
48  _DUMP_FILE = 'capture.pcap'
49
50  def __init__(self, output_path):
51    if not os.path.exists(output_path):
52      os.makedirs(output_path)
53    self._dump_file = os.path.join(output_path, self._DUMP_FILE)
54    self._tmp_output_file = tempfile.NamedTemporaryFile('w', 0)
55    try:
56      self._proc = subprocess.Popen(
57          ['tcpdump'] + _TCP_DUMP_BASE_OPTS + [self._dump_file],
58          stdout=self._tmp_output_file, stderr=subprocess.STDOUT)
59    except OSError as e:
60      raise Exception('Unable to execute TCPDump, please check your '
61          'installation. ' + str(e))
62
63  def CollectProfile(self):
64    self._proc.send_signal(signal.SIGINT)
65    exit_code = self._proc.wait()
66    try:
67      if exit_code:
68        raise Exception(
69            'tcpdump failed with exit code %d. Output:\n%s' %
70            (exit_code, self._GetStdOut()))
71    finally:
72      self._tmp_output_file.close()
73    print 'TCP dump available at: ', self._dump_file
74    print 'Use Wireshark to open it.'
75
76  def _GetStdOut(self):
77    self._tmp_output_file.flush()
78    try:
79      with open(self._tmp_output_file.name) as f:
80        return f.read()
81    except IOError:
82      return ''
83
84
85class TCPDumpProfiler(profiler.Profiler):
86  """A Factory to instantiate the platform-specific profiler."""
87  def __init__(self, browser_backend, platform_backend, output_path):
88    super(TCPDumpProfiler, self).__init__(
89        browser_backend, platform_backend, output_path)
90    if platform_backend.GetOSName() == 'android':
91      self._platform_profiler = _TCPDumpProfilerAndroid(
92          browser_backend.adb, output_path)
93    else:
94      self._platform_profiler = _TCPDumpProfilerLinux(output_path)
95
96  @classmethod
97  def name(cls):
98    return 'tcpdump'
99
100  @classmethod
101  def is_supported(cls, options):
102    if options and options.browser_type.startswith('cros'):
103      return False
104    if sys.platform.startswith('linux'):
105      return True
106    if not options:
107      return android_browser_finder.CanFindAvailableBrowsers()
108    return options.browser_type.startswith('android')
109
110  def CollectProfile(self):
111    self._platform_profiler.CollectProfile()
112