# Copyright 2018 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import json
import logging
import os
import re
import select
import socket
import sys
import subprocess
import tempfile
import time

DIR_SOURCE_ROOT = os.path.abspath(
    os.path.join(os.path.dirname(__file__), os.pardir, os.pardir))
sys.path.append(os.path.join(DIR_SOURCE_ROOT, 'build', 'util', 'lib', 'common'))
import chrome_test_server_spawner

PORT_MAP_RE = re.compile('Allocated port (?P<port>\d+) for remote')
GET_PORT_NUM_TIMEOUT_SECS = 5


def _ConnectPortForwardingTask(target, local_port):
  """Establishes a port forwarding SSH task to a localhost TCP endpoint hosted
  at port |local_port|. Blocks until port forwarding is established.

  Returns the remote port number."""

  forwarding_flags = ['-O', 'forward',  # Send SSH mux control signal.
                      '-R', '0:localhost:%d' % local_port,
                      '-v',   # Get forwarded port info from stderr.
                      '-NT']  # Don't execute command; don't allocate terminal.
  task = target.RunCommandPiped([],
                                ssh_args=forwarding_flags,
                                stderr=subprocess.PIPE)

  # SSH reports the remote dynamic port number over stderr.
  # Unfortunately, the output is incompatible with Python's line buffered
  # input (or vice versa), so we have to build our own buffered input system to
  # pull bytes over the pipe.
  poll_obj = select.poll()
  poll_obj.register(task.stderr, select.POLLIN)
  line = ''
  timeout = time.time() + GET_PORT_NUM_TIMEOUT_SECS
  while time.time() < timeout:
    poll_result = poll_obj.poll(max(0, timeout - time.time()))
    if poll_result:
      next_char = task.stderr.read(1)
      if not next_char:
        break
      line += next_char
      if line.endswith('\n'):
        line = line[:-1]
        logging.debug('ssh: ' + line)
        matched = PORT_MAP_RE.match(line)
        if matched:
          device_port = int(matched.group('port'))
          logging.debug('Port forwarding established (local=%d, device=%d)' %
                        (local_port, device_port))
          task.wait()
          return device_port
        line = ''

  raise Exception('Could not establish a port forwarding connection.')


# Implementation of chrome_test_server_spawner.PortForwarder that uses SSH's
# remote port forwarding feature to forward ports.
class SSHPortForwarder(chrome_test_server_spawner.PortForwarder):
  def __init__(self, target):
    self._target = target

    # Maps the host (server) port to the device port number.
    self._port_mapping = {}

  def Map(self, port_pairs):
    for p in port_pairs:
      _, host_port = p
      self._port_mapping[host_port] = \
          _ConnectPortForwardingTask(self._target, host_port)

  def GetDevicePortForHostPort(self, host_port):
    return self._port_mapping[host_port]

  def Unmap(self, device_port):
    for host_port, entry in self._port_mapping.iteritems():
      if entry == device_port:
        forwarding_args = [
            '-NT', '-O', 'cancel', '-R',
            '%d:localhost:%d' % (self._port_mapping[host_port], host_port)]
        task = self._target.RunCommandPiped([],
                                            ssh_args=forwarding_args,
                                            stderr=subprocess.PIPE)
        task.wait()
        if task.returncode != 0:
          raise Exception(
              'Error %d when unmapping port %d' % (task.returncode,
                                                   device_port))
        del self._port_mapping[host_port]
        return

    raise Exception('Unmap called for unknown port: %d' % device_port)


def SetupTestServer(target, test_concurrency):
  """Provisions a forwarding test server and configures |target| to use it.

  Returns a Popen object for the test server process."""

  logging.debug('Starting test server.')
  spawning_server = chrome_test_server_spawner.SpawningServer(
      0, SSHPortForwarder(target), test_concurrency)
  forwarded_port = _ConnectPortForwardingTask(
      target, spawning_server.server_port)
  spawning_server.Start()

  logging.debug('Test server listening for connections (port=%d)' %
                spawning_server.server_port)
  logging.debug('Forwarded port is %d' % forwarded_port)

  config_file = tempfile.NamedTemporaryFile(delete=True)

  # Clean up the config JSON to only pass ports. See https://crbug.com/810209 .
  config_file.write(json.dumps({
    'name': 'testserver',
    'address': '127.0.0.1',
    'spawner_url_base': 'http://localhost:%d' % forwarded_port
  }))

  config_file.flush()
  target.PutFile(config_file.name, '/data/net-test-server-config')

  return spawning_server