# Copyright (c) 2015 The Chromium OS 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 httplib
import logging
import socket
import tempfile
import time
import xmlrpclib

import common
from autotest_lib.client.bin import utils
from autotest_lib.client.common_lib import error
from autotest_lib.client.common_lib.cros import retry

try:
    import jsonrpclib
except ImportError:
    jsonrpclib = None


class RpcServerTracker(object):
    """
    This class keeps track of all the RPC server connections started on a remote
    host. The caller can use either |xmlrpc_connect| or |jsonrpc_connect| to
    start the required type of rpc server on the remote host.
    The host will cleanup all the open RPC server connections on disconnect.
    """

    _RPC_PROXY_URL_FORMAT = 'http://localhost:%d'
    _RPC_HOST_ADDRESS_FORMAT = 'localhost:%d'
    _RPC_SHUTDOWN_POLLING_PERIOD_SECONDS = 2
    _RPC_SHUTDOWN_TIMEOUT_SECONDS = 10

    def __init__(self, host):
        """
        @param port: The host object associated with this instance of
                     RpcServerTracker.
        """
        self._host = host
        self._rpc_proxy_map = {}


    def _setup_port(self, port, command_name, remote_pid=None):
        """Sets up a tunnel process and register it to rpc_server_tracker.

        Chrome OS on the target closes down most external ports for security.
        We could open the port, but doing that would conflict with security
        tests that check that only expected ports are open.  So, to get to
        the port on the target we use an ssh tunnel.

        This method assumes that xmlrpc and jsonrpc never conflict, since
        we can only either have an xmlrpc or a jsonrpc server listening on
        a remote port. As such, it enforces a single proxy->remote port
        policy, i.e if one starts a jsonrpc proxy/server from port A->B,
        and then tries to start an xmlrpc proxy forwarded to the same port,
        the xmlrpc proxy will override the jsonrpc tunnel process, however:

        1. None of the methods on the xmlrpc proxy will work because
        the server listening on B is jsonrpc.

        2. The xmlrpc client cannot initiate a termination of the JsonRPC
        server, as the only use case currently is goofy, which is tied to
        the factory image. It is much easier to handle a failed xmlrpc
        call on the client than it is to terminate goofy in this scenario,
        as doing the latter might leave the DUT in a hard to recover state.

        With the current implementation newer rpc proxy connections will
        terminate the tunnel processes of older rpc connections tunneling
        to the same remote port. If methods are invoked on the client
        after this has happened they will fail with connection closed errors.

        @param port: The remote forwarding port.
        @param command_name: The name of the remote process, to terminate
                                using pkill.
        @param remote_pid: The PID of the remote background process
                            as a string.

        @return the local port which is used for port forwarding on the ssh
                    client.
        """
        self.disconnect(port)
        local_port = utils.get_unused_port()
        tunnel_proc = self._host.create_ssh_tunnel(port, local_port)
        self._rpc_proxy_map[port] = (command_name, tunnel_proc, remote_pid)
        return local_port


    def _setup_rpc(self, port, command_name, remote_pid=None):
        """Construct a URL for an rpc connection using ssh tunnel.

        @param port: The remote forwarding port.
        @param command_name: The name of the remote process, to terminate
                              using pkill.
        @param remote_pid: The PID of the remote background process
                            as a string.

        @return a url that we can use to initiate the rpc connection.
        """
        return self._RPC_PROXY_URL_FORMAT % self._setup_port(
                port, command_name, remote_pid=remote_pid)


    def tunnel_connect(self, port):
        """Construct a host address using ssh tunnel.

        @param port: The remote forwarding port.

        @return a host address using ssh tunnel.
        """
        return self._RPC_HOST_ADDRESS_FORMAT % self._setup_port(port, None)


    def xmlrpc_connect(self, command, port, command_name=None,
                       ready_test_name=None, timeout_seconds=10,
                       logfile=None, request_timeout_seconds=None):
        """Connect to an XMLRPC server on the host.

        The `command` argument should be a simple shell command that
        starts an XMLRPC server on the given `port`.  The command
        must not daemonize, and must terminate cleanly on SIGTERM.
        The command is started in the background on the host, and a
        local XMLRPC client for the server is created and returned
        to the caller.

        Note that the process of creating an XMLRPC client makes no
        attempt to connect to the remote server; the caller is
        responsible for determining whether the server is running
        correctly, and is ready to serve requests.

        Optionally, the caller can pass ready_test_name, a string
        containing the name of a method to call on the proxy.  This
        method should take no parameters and return successfully only
        when the server is ready to process client requests.  When
        ready_test_name is set, xmlrpc_connect will block until the
        proxy is ready, and throw a TestError if the server isn't
        ready by timeout_seconds.

        If a server is already running on the remote port, this
        method will kill it and disconnect the tunnel process
        associated with the connection before establishing a new one,
        by consulting the rpc_proxy_map in disconnect.

        @param command Shell command to start the server.
        @param port Port number on which the server is expected to
                    be serving.
        @param command_name String to use as input to `pkill` to
            terminate the XMLRPC server on the host.
        @param ready_test_name String containing the name of a
            method defined on the XMLRPC server.
        @param timeout_seconds Number of seconds to wait
            for the server to become 'ready.'  Will throw a
            TestFail error if server is not ready in time.
        @param logfile Logfile to send output when running
            'command' argument.
        @param request_timeout_seconds Timeout in seconds for an XMLRPC request.

        """
        # Clean up any existing state.  If the caller is willing
        # to believe their server is down, we ought to clean up
        # any tunnels we might have sitting around.
        self.disconnect(port)
        remote_pid = None
        if command is not None:
            if logfile:
                remote_cmd = '%s > %s 2>&1' % (command, logfile)
            else:
                remote_cmd = command
            remote_pid = self._host.run_background(remote_cmd)
            logging.debug('Started XMLRPC server on host %s, pid = %s',
                        self._host.hostname, remote_pid)

        # Tunnel through SSH to be able to reach that remote port.
        rpc_url = self._setup_rpc(port, command_name, remote_pid=remote_pid)
        if request_timeout_seconds is not None:
            proxy = TimeoutXMLRPCServerProxy(
                    rpc_url, timeout=request_timeout_seconds, allow_none=True)
        else:
            proxy = xmlrpclib.ServerProxy(rpc_url, allow_none=True)

        if ready_test_name is not None:
            # retry.retry logs each attempt; calculate delay_sec to
            # keep log spam to a dull roar.
            @retry.retry((socket.error,
                          xmlrpclib.ProtocolError,
                          httplib.BadStatusLine),
                         timeout_min=timeout_seconds / 60.0,
                         delay_sec=min(max(timeout_seconds / 20.0, 0.1), 1))
            def ready_test():
                """ Call proxy.ready_test_name(). """
                getattr(proxy, ready_test_name)()
            successful = False
            try:
                logging.info('Waiting %d seconds for XMLRPC server '
                             'to start.', timeout_seconds)
                ready_test()
                successful = True
            finally:
                if not successful:
                    logging.error('Failed to start XMLRPC server.')
                    if logfile:
                        with tempfile.NamedTemporaryFile() as temp:
                            self._host.get_file(logfile, temp.name)
                            logging.error('The log of XML RPC server:\n%s',
                                          open(temp.name).read())
                    self.disconnect(port)
        logging.info('XMLRPC server started successfully.')
        return proxy


    def jsonrpc_connect(self, port):
        """Creates a jsonrpc proxy connection through an ssh tunnel.

        This method exists to facilitate communication with goofy (which is
        the default system manager on all factory images) and as such, leaves
        most of the rpc server sanity checking to the caller. Unlike
        xmlrpc_connect, this method does not facilitate the creation of a remote
        jsonrpc server, as the only clients of this code are factory tests,
        for which the goofy system manager is built in to the image and starts
        when the target boots.

        One can theoretically create multiple jsonrpc proxies all forwarded
        to the same remote port, provided the remote port has an rpc server
        listening. However, in doing so we stand the risk of leaking an
        existing tunnel process, so we always disconnect any older tunnels
        we might have through disconnect.

        @param port: port on the remote host that is serving this proxy.

        @return: The client proxy.
        """
        if not jsonrpclib:
            logging.warning('Jsonrpclib could not be imported. Check that '
                            'site-packages contains jsonrpclib.')
            return None

        proxy = jsonrpclib.jsonrpc.ServerProxy(self._setup_rpc(port, None))

        logging.info('Established a jsonrpc connection through port %s.', port)
        return proxy


    def disconnect(self, port):
        """Disconnect from an RPC server on the host.

        Terminates the remote RPC server previously started for
        the given `port`.  Also closes the local ssh tunnel created
        for the connection to the host.  This function does not
        directly alter the state of a previously returned RPC
        client object; however disconnection will cause all
        subsequent calls to methods on the object to fail.

        This function does nothing if requested to disconnect a port
        that was not previously connected via _setup_rpc.

        @param port Port number passed to a previous call to
                    `_setup_rpc()`.
        """
        if port not in self._rpc_proxy_map:
            return
        remote_name, tunnel_proc, remote_pid = self._rpc_proxy_map[port]
        if remote_name:
            # We use 'pkill' to find our target process rather than
            # a PID, because the host may have rebooted since
            # connecting, and we don't want to kill an innocent
            # process with the same PID.
            #
            # 'pkill' helpfully exits with status 1 if no target
            # process  is found, for which run() will throw an
            # exception.  We don't want that, so we the ignore
            # status.
            self._host.run("pkill -f '%s'" % remote_name, ignore_status=True)
            if remote_pid:
                logging.info('Waiting for RPC server "%s" shutdown',
                             remote_name)
                start_time = time.time()
                while (time.time() - start_time <
                       self._RPC_SHUTDOWN_TIMEOUT_SECONDS):
                    running_processes = self._host.run(
                            "pgrep -f '%s'" % remote_name,
                            ignore_status=True).stdout.split()
                    if not remote_pid in running_processes:
                        logging.info('Shut down RPC server.')
                        break
                    time.sleep(self._RPC_SHUTDOWN_POLLING_PERIOD_SECONDS)
                else:
                    raise error.TestError('Failed to shutdown RPC server %s' %
                                          remote_name)

        self._host.disconnect_ssh_tunnel(tunnel_proc, port)
        del self._rpc_proxy_map[port]


    def disconnect_all(self):
        """Disconnect all known RPC proxy ports."""
        for port in self._rpc_proxy_map.keys():
            self.disconnect(port)


class TimeoutXMLRPCServerProxy(xmlrpclib.ServerProxy):
    """XMLRPC ServerProxy supporting timeout."""
    def __init__(self, uri, timeout=20, *args, **kwargs):
        """Initializes a TimeoutXMLRPCServerProxy.

        @param uri: URI to a XMLRPC server.
        @param timeout: Timeout in seconds for a XMLRPC request.
        @param *args: args to xmlrpclib.ServerProxy.
        @param **kwargs: kwargs to xmlrpclib.ServerProxy.

        """
        if timeout:
            kwargs['transport'] = TimeoutXMLRPCTransport(timeout=timeout)
        xmlrpclib.ServerProxy.__init__(self, uri, *args, **kwargs)


class TimeoutXMLRPCTransport(xmlrpclib.Transport):
    """A Transport subclass supporting timeout."""
    def __init__(self, timeout=20, *args, **kwargs):
        """Initializes a TimeoutXMLRPCTransport.

        @param timeout: Timeout in seconds for a HTTP request through this transport layer.
        @param *args: args to xmlrpclib.Transport.
        @param **kwargs: kwargs to xmlrpclib.Transport.

        """
        xmlrpclib.Transport.__init__(self, *args, **kwargs)
        self.timeout = timeout


    def make_connection(self, host):
        """Overwrites make_connection in xmlrpclib.Transport with timeout.

        @param host: Host address to connect.

        @return: A httplib.HTTPConnection connecting to host with timeout.

        """
        conn = httplib.HTTPConnection(host, timeout=self.timeout)
        return conn