1# Copyright 2016 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
5"""Start and stop tsproxy."""
6
7import logging
8import os
9import re
10import subprocess
11import sys
12
13from telemetry.core import util
14from telemetry.internal.util import atexit_with_log
15
16import py_utils
17
18
19_TSPROXY_PATH = os.path.join(
20    util.GetTelemetryThirdPartyDir(), 'tsproxy', 'tsproxy.py')
21
22
23def ParseTsProxyPortFromOutput(output_line):
24  port_re = re.compile(
25      r'Started Socks5 proxy server on '
26      r'(?P<host>[^:]*):'
27      r'(?P<port>\d+)')
28  m = port_re.match(output_line.strip())
29  if m:
30    return int(m.group('port'))
31
32
33class TsProxyServer(object):
34  """Start and Stop Tsproxy.
35
36  TsProxy provides basic latency, download and upload traffic shaping. This
37  class provides a programming API to the tsproxy script in
38  telemetry/third_party/tsproxy/tsproxy.py
39  """
40
41  def __init__(self, host_ip=None, http_port=None, https_port=None):
42    """Initialize TsProxyServer.
43    """
44    self._proc = None
45    self._port = None
46    self._is_running = False
47    self._host_ip = host_ip
48    assert bool(http_port) == bool(https_port)
49    self._http_port = http_port
50    self._https_port = https_port
51
52  @property
53  def port(self):
54    return self._port
55
56  def StartServer(self, timeout=10):
57    """Start TsProxy server and verify that it started.
58    """
59    cmd_line = [sys.executable, _TSPROXY_PATH]
60    cmd_line.extend([
61        '--port=0'])  # Use port 0 so tsproxy picks a random available port.
62    if self._host_ip:
63      cmd_line.append('--desthost=%s' % self._host_ip)
64    if self._http_port:
65      cmd_line.append(
66        '--mapports=443:%s,*:%s' % (self._https_port, self._http_port))
67    logging.info('Tsproxy commandline: %r' % cmd_line)
68    self._proc = subprocess.Popen(
69        cmd_line, stdout=subprocess.PIPE, stdin=subprocess.PIPE,
70        stderr=subprocess.PIPE, bufsize=1)
71    atexit_with_log.Register(self.StopServer)
72    try:
73      py_utils.WaitFor(self._IsStarted, timeout)
74      logging.info('TsProxy port: %s', self._port)
75      self._is_running = True
76    except py_utils.TimeoutException:
77      err = self.StopServer()
78      raise RuntimeError(
79          'Error starting tsproxy: %s' % err)
80
81  def _IsStarted(self):
82    assert not self._is_running
83    assert self._proc
84    if self._proc.poll() is not None:
85      return False
86    self._proc.stdout.flush()
87    self._port = ParseTsProxyPortFromOutput(
88          output_line=self._proc.stdout.readline())
89    return self._port != None
90
91
92  def _IssueCommand(self, command_string, timeout):
93    logging.info('Issuing command to ts_proxy_server: %s', command_string)
94    command_output = []
95    self._proc.stdin.write('%s\n' % command_string)
96    self._proc.stdin.flush()
97    self._proc.stdout.flush()
98    def CommandStatusIsRead():
99      command_output.append(self._proc.stdout.readline().strip())
100      return (
101          command_output[-1] == 'OK' or command_output[-1] == 'ERROR')
102    py_utils.WaitFor(CommandStatusIsRead, timeout)
103    if not 'OK' in command_output:
104      raise RuntimeError('Failed to execute command %s:\n%s' %
105                         (repr(command_string), '\n'.join(command_output)))
106
107
108  def UpdateOutboundPorts(self, http_port, https_port, timeout=5):
109    assert http_port and https_port
110    assert http_port != https_port
111    assert isinstance(http_port, int) and isinstance(https_port, int)
112    assert 1 <= http_port <= 65535
113    assert 1 <= https_port <= 65535
114    self._IssueCommand('set mapports 443:%i,*:%i' % (https_port, http_port),
115                       timeout)
116
117  def UpdateTrafficSettings(self, round_trip_latency_ms=0,
118      download_bandwidth_kbps=0, upload_bandwidth_kbps=0, timeout=5):
119    self._IssueCommand('set rtt %s' % round_trip_latency_ms, timeout)
120    self._IssueCommand('set inkbps %s' % download_bandwidth_kbps, timeout)
121    self._IssueCommand('set outkbps %s' % upload_bandwidth_kbps, timeout)
122
123  def StopServer(self):
124    """Stop TsProxy Server."""
125    if not self._is_running:
126      logging.debug('Attempting to stop TsProxy server that is not running.')
127      return
128    if self._proc:
129      self._proc.terminate()
130      self._proc.wait()
131    err = self._proc.stderr.read()
132    self._proc = None
133    self._port = None
134    self._is_running = False
135    return err
136
137  def __enter__(self):
138    """Add support for with-statement."""
139    self.StartServer()
140    return self
141
142  def __exit__(self, unused_exc_type, unused_exc_val, unused_exc_tb):
143    """Add support for with-statement."""
144    self.StopServer()
145