• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
7This helper will help with communicating to a SCPI device which is behind
8a firewall and can be accessed through a SSH tunnel. Please make sure you
9can login to SSH Tunnel machine without a password by configuring the ssh
10pub 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
29import logging
30import shlex
31import socket
32import subprocess
33import sys
34import time
35
36from autotest_lib.client.common_lib import utils
37from autotest_lib.server.cros.network.rf_switch import scpi
38
39
40class 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