server/cros/network/rf_switch/scpi_ssh_tunnel.py

# Copyright (c) 2017 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.

"""Interface for SCPI Protocol through a SSH tunnel.

This helper will help with communicating to a SCPI device which is behind
a firewall and can be accessed through a SSH tunnel. Please make sure you
can login to SSH Tunnel machine without a password by configuring the ssh
pub keys.

                       /|
                      / |
                     /  |
                    |   |
    +---------+     |   |     +--------+      +--------+
    |         |     |   |     |        |      |        |
    |  Test   +-----|  -------+  SSH   +------+  SCPI  |
    | Machine |     |   |     | Tunnel |      | Device |
    |         |     |   |     |        |      |        |
    +---------+     |   |     +--------+      +--------+
                    |  /
                    | /
                    |/

                 Firewall
"""

from __future__ import print_function

import logging
import shlex
import six
import socket
import subprocess
import sys
import time

from autotest_lib.client.common_lib import utils
from autotest_lib.server.cros.network.rf_switch import scpi


class ScpiSshTunnel(scpi.Scpi):
    """Class for SCPI protocol though SSH tunnel."""

    _TUNNEL_ESTABLISH_TIME_SECS = 10

    def __init__(self,
                 host,
                 proxy_host,
                 proxy_username,
                 port=scpi.Scpi.SCPI_PORT,
                 proxy_port=None):
        """SCPI handler through a proxy server.

        @param host: Hostname or IP address of SCPI device
        @param proxy_host: Hostname or IP address of SSH tunnel device
        @param proxy_username: Username for SSH tunnel device
        @param port: SCPI port on device (default 5025)
        @param proxy_port: port number to bind for SSH tunnel. If
            none will pick an available free port

        """
        self.host = host
        self.port = port
        self.proxy_host = proxy_host
        self.proxy_username = proxy_username
        self.proxy_port = proxy_port or utils.get_unused_port()

        # We will override the parent initialize method and use a tunnel
        # to connect to the socket connection

        # Start SSH tunnel
        try:
            tunnel_command = self._make_tunnel_command()
            logging.debug('Tunnel command: %s', tunnel_command)
            args = shlex.split(tunnel_command)
            self.ssh_tunnel = subprocess.Popen(args)
            time.sleep(self._TUNNEL_ESTABLISH_TIME_SECS)
            logging.debug(
                'Started ssh tunnel, local = %d, remote = %d, pid = %d',
                self.proxy_port, self.port, self.ssh_tunnel.pid)
        except OSError as e:
            logging.exception('Error starting SSH tunnel to SCPI device.')
            six.reraise(scpi.ScpiException(cause=e), None, sys.exc_info()[2])

        # Open a socket connection for communication with chassis
        # using the SSH Tunnel.
        try:
            self.socket = socket.socket()
            self.socket.connect(('127.0.0.1', self.proxy_port))
        except (socket.error, socket.timeout) as e:
            logging.error('Error connecting to SCPI device.')
            six.reraise(scpi.ScpiException(cause=e), None, sys.exc_info()[2])

    def _make_tunnel_command(self, hosts_file='/dev/null',
                             connect_timeout=30, alive_interval=300):
        tunnel_command = ('/usr/bin/ssh -ax -o StrictHostKeyChecking=no'
                          ' -o UserKnownHostsFile=%s -o BatchMode=yes'
                          ' -o ConnectTimeout=%s -o ServerAliveInterval=%s'
                          ' -l %s -L %s:%s:%s  -nNq %s') % (
                              hosts_file, connect_timeout, alive_interval,
                              self.proxy_username, self.proxy_port,
                              self.host, self.port, self.proxy_host)
        return tunnel_command

    def close(self):
        """Close the connection."""
        if hasattr(self, 's'):
            self.socket.close()
            del self.socket
        if self.ssh_tunnel:
            self.ssh_tunnel.kill()
            del self.ssh_tunnel