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