• 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
29from __future__ import print_function
30
31import logging
32import shlex
33import six
34import socket
35import subprocess
36import sys
37import time
38
39from autotest_lib.client.common_lib import utils
40from autotest_lib.server.cros.network.rf_switch import scpi
41
42
43class ScpiSshTunnel(scpi.Scpi):
44    """Class for SCPI protocol though SSH tunnel."""
45
46    _TUNNEL_ESTABLISH_TIME_SECS = 10
47
48    def __init__(self,
49                 host,
50                 proxy_host,
51                 proxy_username,
52                 port=scpi.Scpi.SCPI_PORT,
53                 proxy_port=None):
54        """SCPI handler through a proxy server.
55
56        @param host: Hostname or IP address of SCPI device
57        @param proxy_host: Hostname or IP address of SSH tunnel device
58        @param proxy_username: Username for SSH tunnel device
59        @param port: SCPI port on device (default 5025)
60        @param proxy_port: port number to bind for SSH tunnel. If
61            none will pick an available free port
62
63        """
64        self.host = host
65        self.port = port
66        self.proxy_host = proxy_host
67        self.proxy_username = proxy_username
68        self.proxy_port = proxy_port or utils.get_unused_port()
69
70        # We will override the parent initialize method and use a tunnel
71        # to connect to the socket connection
72
73        # Start SSH tunnel
74        try:
75            tunnel_command = self._make_tunnel_command()
76            logging.debug('Tunnel command: %s', tunnel_command)
77            args = shlex.split(tunnel_command)
78            self.ssh_tunnel = subprocess.Popen(args)
79            time.sleep(self._TUNNEL_ESTABLISH_TIME_SECS)
80            logging.debug(
81                'Started ssh tunnel, local = %d, remote = %d, pid = %d',
82                self.proxy_port, self.port, self.ssh_tunnel.pid)
83        except OSError as e:
84            logging.exception('Error starting SSH tunnel to SCPI device.')
85            six.reraise(scpi.ScpiException(cause=e), None, sys.exc_info()[2])
86
87        # Open a socket connection for communication with chassis
88        # using the SSH Tunnel.
89        try:
90            self.socket = socket.socket()
91            self.socket.connect(('127.0.0.1', self.proxy_port))
92        except (socket.error, socket.timeout) as e:
93            logging.error('Error connecting to SCPI device.')
94            six.reraise(scpi.ScpiException(cause=e), None, sys.exc_info()[2])
95
96    def _make_tunnel_command(self, hosts_file='/dev/null',
97                             connect_timeout=30, alive_interval=300):
98        tunnel_command = ('/usr/bin/ssh -ax -o StrictHostKeyChecking=no'
99                          ' -o UserKnownHostsFile=%s -o BatchMode=yes'
100                          ' -o ConnectTimeout=%s -o ServerAliveInterval=%s'
101                          ' -l %s -L %s:%s:%s  -nNq %s') % (
102                              hosts_file, connect_timeout, alive_interval,
103                              self.proxy_username, self.proxy_port,
104                              self.host, self.port, self.proxy_host)
105        return tunnel_command
106
107    def close(self):
108        """Close the connection."""
109        if hasattr(self, 's'):
110            self.socket.close()
111            del self.socket
112        if self.ssh_tunnel:
113            self.ssh_tunnel.kill()
114            del self.ssh_tunnel
115