Home | History | Annotate | Download | only in rf_switch 
      1 # Copyright (c) 2017 The Chromium OS 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 """Interface for SCPI Protocol through a SSH tunnel.
      6 
      7 This helper will help with communicating to a SCPI device which is behind
      8 a firewall and can be accessed through a SSH tunnel. Please make sure you
      9 can login to SSH Tunnel machine without a password by configuring the ssh
     10 pub keys.
     11 
     12                        /|
     13                       / |
     14                      /  |
     15                     |   |
     16     +---------+     |   |     +--------+      +--------+
     17     |         |     |   |     |        |      |        |
     18     |  Test   +-----|  -------+  SSH   +------+  SCPI  |
     19     | Machine |     |   |     | Tunnel |      | Device |
     20     |         |     |   |     |        |      |        |
     21     +---------+     |   |     +--------+      +--------+
     22                     |  /
     23                     | /
     24                     |/
     25 
     26                  Firewall
     27 """
     28 
     29 import logging
     30 import shlex
     31 import socket
     32 import subprocess
     33 import sys
     34 import time
     35 
     36 from autotest_lib.client.common_lib import utils
     37 from autotest_lib.server.cros.network.rf_switch import scpi
     38 
     39 
     40 class ScpiSshTunnel(scpi.Scpi):
     41     """Class for SCPI protocol though SSH tunnel."""
     42 
     43     _TUNNEL_ESTABLISH_TIME_SECS = 10
     44 
     45     def __init__(self,
     46                  host,
     47                  proxy_host,
     48                  proxy_username,
     49                  port=scpi.Scpi.SCPI_PORT,
     50                  proxy_port=None):
     51         """SCPI handler through a proxy server.
     52 
     53         @param host: Hostname or IP address of SCPI device
     54         @param proxy_host: Hostname or IP address of SSH tunnel device
     55         @param proxy_username: Username for SSH tunnel device
     56         @param port: SCPI port on device (default 5025)
     57         @param proxy_port: port number to bind for SSH tunnel. If
     58             none will pick an available free port
     59 
     60         """
     61         self.host = host
     62         self.port = port
     63         self.proxy_host = proxy_host
     64         self.proxy_username = proxy_username
     65         self.proxy_port = proxy_port or utils.get_unused_port()
     66 
     67         # We will override the parent initialize method and use a tunnel
     68         # to connect to the socket connection
     69 
     70         # Start SSH tunnel
     71         try:
     72             tunnel_command = self._make_tunnel_command()
     73             logging.debug('Tunnel command: %s', tunnel_command)
     74             args = shlex.split(tunnel_command)
     75             self.ssh_tunnel = subprocess.Popen(args)
     76             time.sleep(self._TUNNEL_ESTABLISH_TIME_SECS)
     77             logging.debug(
     78                 'Started ssh tunnel, local = %d, remote = %d, pid = %d',
     79                 self.proxy_port, self.port, self.ssh_tunnel.pid)
     80         except OSError as e:
     81             logging.exception('Error starting SSH tunnel to SCPI device.')
     82             raise scpi.ScpiException(cause=e), None, sys.exc_info()[2]
     83 
     84         # Open a socket connection for communication with chassis
     85         # using the SSH Tunnel.
     86         try:
     87             self.socket = socket.socket()
     88             self.socket.connect(('127.0.0.1', self.proxy_port))
     89         except (socket.error, socket.timeout) as e:
     90             logging.error('Error connecting to SCPI device.')
     91             raise scpi.ScpiException(cause=e), None, sys.exc_info()[2]
     92 
     93     def _make_tunnel_command(self, hosts_file='/dev/null',
     94                              connect_timeout=30, alive_interval=300):
     95         tunnel_command = ('/usr/bin/ssh -ax -o StrictHostKeyChecking=no'
     96                           ' -o UserKnownHostsFile=%s -o BatchMode=yes'
     97                           ' -o ConnectTimeout=%s -o ServerAliveInterval=%s'
     98                           ' -l %s -L %s:%s:%s  -nNq %s') % (
     99                               hosts_file, connect_timeout, alive_interval,
    100                               self.proxy_username, self.proxy_port,
    101                               self.host, self.port, self.proxy_host)
    102         return tunnel_command
    103 
    104     def close(self):
    105         """Close the connection."""
    106         if hasattr(self, 's'):
    107             self.socket.close()
    108             del self.socket
    109         if self.ssh_tunnel:
    110             self.ssh_tunnel.kill()
    111             del self.ssh_tunnel
    112