From 790c42e34e194439ee0ce49d9ab5f31adebdfc6b Mon Sep 17 00:00:00 2001 From: Rouven Czerwinski Date: Wed, 4 Apr 2018 11:53:29 +0200 Subject: [PATCH 1/4] multiple: move get_free_ports to labgrid.util Signed-off-by: Rouven Czerwinski --- labgrid/remote/exporter.py | 7 +------ labgrid/util/__init__.py | 2 ++ labgrid/util/ports.py | 8 ++++++++ 3 files changed, 11 insertions(+), 6 deletions(-) create mode 100644 labgrid/util/ports.py diff --git a/labgrid/remote/exporter.py b/labgrid/remote/exporter.py index 211827dad..d5ec82193 100755 --- a/labgrid/remote/exporter.py +++ b/labgrid/remote/exporter.py @@ -15,6 +15,7 @@ from .config import ResourceConfig from .common import ResourceEntry, enable_tcp_nodelay +from ..util import get_free_port try: import pkg_resources @@ -22,12 +23,6 @@ except pkg_resources.DistributionNotFound: __version__ = "unknown" -def get_free_port(): - """Helper function to always return an unused port.""" - with closing(socket(AF_INET, SOCK_STREAM)) as s: - s.bind(('', 0)) - return s.getsockname()[1] - exports = {} reexec = False diff --git a/labgrid/util/__init__.py b/labgrid/util/__init__.py index 53cbd6b05..4953fa4ae 100644 --- a/labgrid/util/__init__.py +++ b/labgrid/util/__init__.py @@ -3,3 +3,5 @@ from .timeout import Timeout from .marker import gen_marker from .yaml import load, dump +from .ssh import SSHManager +from .ports import get_free_port diff --git a/labgrid/util/ports.py b/labgrid/util/ports.py new file mode 100644 index 000000000..fd4536ba6 --- /dev/null +++ b/labgrid/util/ports.py @@ -0,0 +1,8 @@ +import socket +from contextlib import closing + +def get_free_port(()): + """Helper function to always return an unused port.""" + with closing(socket(AF_INET, SOCK_STREAM)) as s: + s.bind(('', 0)) + return s.getsockname()[1] From 1b714b5564508b3146d00d766d1a42d854ccdf58 Mon Sep 17 00:00:00 2001 From: Rouven Czerwinski Date: Wed, 25 Apr 2018 14:55:01 +0200 Subject: [PATCH 2/4] travis: add key for localhost Signed-off-by: Rouven Czerwinski --- .travis.yml | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/.travis.yml b/.travis.yml index f754c1d84..a40639d9c 100644 --- a/.travis.yml +++ b/.travis.yml @@ -1,10 +1,21 @@ language: python + +addons: + apt: + packages: + libow-dev + openssh-server + openssh-client + python: - "3.5" - "3.6" before_install: - - sudo apt-get -qq update - - sudo apt-get install -y libow-dev + - ssh-keygen -f ~/.ssh/localkey -t ed25519 -N "" + - cat ~/.ssh/localkey.pub > ~/.ssh/authorized_keys + - echo -e "Host localhost\n Identityfile ~/.ssh/localkey" > ~/.ssh/config + - cat ~/.ssh/config + - ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null localhost -- echo "Hello" install: - pip install -r travis-requirements.txt - pip install -e . From 73f29e1e1a7796064c7cbda71929e1bdb1639259 Mon Sep 17 00:00:00 2001 From: Rouven Czerwinski Date: Sun, 11 Mar 2018 22:27:49 +0100 Subject: [PATCH 3/4] util/ssh: add SSHConnection, SSHConnectionManager and SSHMANAGER The new SSHMANAGER allows us to manage SSHConnections to remote hosts for the distributed infrastructure. It supports port forwards and file uploads, making only a ssh connection to the remote host necessary. The SSHMANAGER is exported for use by other modules. Signed-off-by: Rouven Czerwinski --- labgrid/util/__init__.py | 1 - labgrid/util/ports.py | 4 +- labgrid/util/ssh.py | 296 +++++++++++++++++++++++++++++++++++++++ tests/conftest.py | 8 ++ tests/test_util.py | 128 +++++++++++++++++ 5 files changed, 434 insertions(+), 3 deletions(-) create mode 100644 labgrid/util/ssh.py diff --git a/labgrid/util/__init__.py b/labgrid/util/__init__.py index 4953fa4ae..29b12de92 100644 --- a/labgrid/util/__init__.py +++ b/labgrid/util/__init__.py @@ -3,5 +3,4 @@ from .timeout import Timeout from .marker import gen_marker from .yaml import load, dump -from .ssh import SSHManager from .ports import get_free_port diff --git a/labgrid/util/ports.py b/labgrid/util/ports.py index fd4536ba6..d6b26f3b5 100644 --- a/labgrid/util/ports.py +++ b/labgrid/util/ports.py @@ -1,7 +1,7 @@ -import socket +from socket import gethostname, socket, AF_INET, SOCK_STREAM from contextlib import closing -def get_free_port(()): +def get_free_port(): """Helper function to always return an unused port.""" with closing(socket(AF_INET, SOCK_STREAM)) as s: s.bind(('', 0)) diff --git a/labgrid/util/ssh.py b/labgrid/util/ssh.py new file mode 100644 index 000000000..9869bbb06 --- /dev/null +++ b/labgrid/util/ssh.py @@ -0,0 +1,296 @@ +# pylint: disable=no-member +import tempfile +import logging +import subprocess +import os +from functools import wraps + +import attr +from ..driver.exception import ExecutionError + +from .ports import get_free_port + +__all__ = ['SSHMANAGER', 'SSHConnection', 'ForwardError'] + + +@attr.s +class SSHConnectionManager: + _connections = attr.ib( + default=attr.Factory(dict), + validator=attr.validators.optional(attr.validators.instance_of(dict)) + ) + _tmpdir = attr.ib( + default=attr. + Factory(lambda: tempfile.mkdtemp(prefix='labgrid-ssh-manager-')), + validator=attr.validators.optional(attr.validators.instance_of(str)) + ) + + def __attrs_post_init__(self): + self.logger = logging.getLogger("{}".format(self)) + + def get(self, host: str): + instance = self._connections.get(host) + if instance is None: + # pylint: disable=unsupported-assignment-operation + self.logger.debug("Trying to start new control socket") + # instance = self._start_control_socket(host) + instance = SSHConnection(host) + instance.connect() + self._connections[host] = instance + return instance + + def add_connection(self, connection): + # pylint: disable=unsupported-assignment-operation + assert isinstance(connection, SSHConnection) + if connection.host not in self._connections: + self._connections[connection.host] = connection + + def remove_connection(self, connection): + # pylint: disable=unsupported-assignment-operation + assert isinstance(connection, SSHConnection) + if connection.isactive(): + raise ExecutionError("Can't remove active connection") + self._connections[connection.host] = connection + + def open(self, host): + con = self.get(host) + return con + + def close(self, host): + con = self.get(host) + con.disconnect() + self.remove_connection(con) + + def request_forward(self, host, port): + con = self.get(host) + return con.add_port_forward(port) + + def remove_forward(self, host, port): + con = self.get(host) + con.remove_port_forward(port) + + def put_file(self, host, local_file, dest_file): + con = self.get(host) + con.put_file(local_file, dest_file) + + def get_file(self, host, dest_file, local_file): + con = self.get(host) + con.get_file(dest_file, local_file) + + +def check_active(func): + """Check if an SSHConnection is active as a decorator""" + + @wraps(func) + def wrapper(cls, *_args, **_kwargs): + if not cls.isactive(): + raise ExecutionError( + "{} can not be called ({} is not active)".format( + func.__qualname__, cls + ) + ) + return func(cls, *_args, **_kwargs) + + return wrapper + + +@attr.s +class SSHConnection: + """SSHConnections are individual connections to hosts managed by a control + socket. In addition to command execution this class also provides an + interface to manage port forwardings. These are used in the remote + infrastructure to tunnel multiple connections over one SSH link. + + A public identity infrastructure is assumed, no extra username or passwords + are supported.""" + host = attr.ib(validator=attr.validators.instance_of(str)) + _active = attr.ib( + default=False, validator=attr.validators.instance_of(bool) + ) + _tmpdir = attr.ib( + default=attr. + Factory(lambda: tempfile.mkdtemp(prefix="labgrid-connection-")), + validator=attr.validators.instance_of(str) + ) + _forwards = attr.ib(default=attr.Factory(dict)) + + def __attrs_post_init__(self): + self._logger = logging.getLogger("{}".format(self)) + self._ssh_prefix = "-o LogLevel=ERROR" + self._ssh_prefix = "-o PasswordAuthentication=no" + self._socket = os.path.join( + self._tmpdir, 'control-{}'.format(self.host) + ) + + def _open_connection(self): + """Internal function which appends the control socket and checks if the + connection is already open""" + self._ssh_prefix += " -o ControlPath={}".format( + self._socket + ) if self._check_master() else "" + + def _run_socket_command(self, command, forward=""): + "Internal function to send a command to the control socket" + complete_cmd = "ssh -x -o ControlPath={cpath} -O {command}{forward} {host}".format( + cpath=self._socket, + command=command, + forward=forward, + host=self.host + ).split(' ') + res = subprocess.check_call( + complete_cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=2 + ) + + return res + + def _run_command(self, command): + "Internal function to run a command over the SSH connection" + complete_cmd = "ssh -x -o ControlPath={cpath} {host} {command}".format( + cpath=self._socket, + host=self.host, + command=command, + ).split(' ') + res = subprocess.check_call( + complete_cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL + ) + + return res + + @check_active + def run_command(self, command): + """Run a command over the SSHConnection + + Args: + command (string): The command to run + + Returns: + int: exitcode of the command + """ + return self._run_command(command) + + @check_active + def get_file(self, remote_file, local_file): + """Get a file from the remote host""" + subprocess.check_call([ + "scp", "-o", "ControlPath={}".format(self._socket), + "{}:{}".format(self.host, remote_file), "{}".format(local_file) + ]) + + @check_active + def put_file(self, local_file, remote_path): + """Put a file onto the remote host""" + subprocess.check_call([ + "scp", "-o", "ControlPath={}".format(self._socket), + "{}".format(local_file), "{}:{}".format(self.host, remote_path) + ]) + + @check_active + def add_port_forward(self, remote_port): + """forward command""" + local_port = get_free_port() + + # pylint: disable=not-an-iterable + if remote_port in self._forwards: + return self._forwards[remote_port] + self._run_socket_command( + "forward", " -L {local}:localhost:{remote}".format( + local=local_port, remote=remote_port + ) + ) + self._forwards[remote_port] = local_port + return local_port + + @check_active + def remove_port_forward(self, remote_port): + """cancel command""" + local_port = self._forwards.pop(remote_port, None) + + # pylint: disable=not-an-iterable + if local_port is None: + raise ForwardError("Forward does not exist") + + self._run_socket_command( + "cancel", " -L {local}:localhost:{remote}".format( + local=local_port, remote=remote_port + ) + ) + + def connect(self): + self._open_connection() + self._active = True + + @check_active + def disconnect(self): + self._disconnect() + + def isactive(self): + return self._active + + def _check_master(self): + args = ["ssh", "-O", "check", "{}".format(self.host)] + check = subprocess.call( + args, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL + ) + if check == 0: + return "" + + return self._start_own_master() + + def _start_own_master(self): + """Starts a controlmaster connection in a temporary directory.""" + control = os.path.join(self._tmpdir, 'control-{}'.format(self.host)) + args = ( + "ssh -n {} -x -o ConnectTimeout=30 -o ControlPersist=300 " + "-o UserKnownHostsFile=/dev/null " + "-o StrictHostKeyChecking=no -MN -S {} {}" + ).format(self._ssh_prefix, control, self.host).split(" ") + + self.process = subprocess.Popen( + args, + stderr=subprocess.PIPE, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE + ) + + try: + if self.process.wait(timeout=30) is not 0: + raise ExecutionError( + "failed to connect to {} with {} and {}".format( + self.host, args, self.process.wait() + ) + ) + except subprocess.TimeoutExpired: + raise ExecutionError( + "failed to connect to {} with {} and {}".format( + self.host, args, self.process.wait() + ) + ) + + if not os.path.exists(control): + raise ExecutionError("no control socket to {}".format(self.host)) + + self._logger.debug('Connected to %s', self.host) + + return control + + def _disconnect(self): + self._run_socket_command("exit") + self._active = False + + +SSHMANAGER = SSHConnectionManager() + + +@attr.s +class ForwardError(Exception): + msg = attr.ib(validator=attr.validators.instance_of(str)) diff --git a/tests/conftest.py b/tests/conftest.py index b94e6ac5f..dfd8869e9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -86,14 +86,22 @@ def exporter(tmpdir): def pytest_addoption(parser): parser.addoption("--sigrok-usb", action="store_true", help="Run sigrok usb tests with fx2lafw device (0925:3881)") + parser.addoption("--local-sshmanager", action="store_true", + help="Run SSHManager tests against localhost") def pytest_configure(config): # register an additional marker config.addinivalue_line("markers", "sigrokusb: enable fx2lafw USB tests (0925:3881)") + config.addinivalue_line("markers", + "localsshmanager: test SSHManager against Localhost") def pytest_runtest_setup(item): envmarker = item.get_marker("sigrokusb") if envmarker is not None: if item.config.getoption("--sigrok-usb") is False: pytest.skip("sigrok usb tests not enabled (enable with --sigrok-usb)") + envmarker = item.get_marker("localsshmanager") + if envmarker is not None: + if item.config.getoption("--local-sshmanager") is False: + pytest.skip("SSHManager tests against localhost not enabled (enable with --local-sshmanager)") diff --git a/tests/test_util.py b/tests/test_util.py index cfbb50002..b87090279 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -1,7 +1,25 @@ +import os.path +import subprocess + import attr import pytest from labgrid.util import diff_dict, flat_dict, filter_dict +from labgrid.util.ssh import ForwardError +from labgrid.util.ssh import SSHConnection +from labgrid.driver.exception import ExecutionError + +@pytest.fixture +def connection_localhost(): + con = SSHConnection("localhost") + con.connect() + yield con + con.disconnect() + +@pytest.fixture +def sshmanager(): + from labgrid.util.ssh import SSHMANAGER + return SSHMANAGER def test_diff_dict(): dict_a = {"a": 1, @@ -41,3 +59,113 @@ class A: assert str(record[1].message) == "unsupported attribute 'baz' with value '3' for class 'A'" assert d_filtered is not d_orig assert d_filtered == {'foo': 1} + +def test_sshmanager_get(): + from labgrid.util.ssh import SSHMANAGER + assert SSHMANAGER != None + +def test_sshconnection_get(): + from labgrid.util.ssh import SSHConnection + SSHConnection("localhost") + +def test_sshconnection_inactive_raise(): + from labgrid.util.ssh import SSHConnection + con = SSHConnection("localhost") + with pytest.raises(ExecutionError): + con.run_command("echo Hallo") + +def test_sshconnection_connect(connection_localhost): + assert connection_localhost.isactive() + assert os.path.exists(connection_localhost._socket) + +def test_sshconnection_run(connection_localhost): + assert connection_localhost.run_command("echo Hello") == 0 + +def test_sshconnection_port_forward_add_remove(connection_localhost): + port = 1337 + test_string = "Hello World" + + local_port = connection_localhost.add_port_forward(port) + nc_listen = subprocess.Popen(["nc", "-l", "-p", "1337"], stdout=subprocess.PIPE, stdin=None) + nc_send = subprocess.Popen(["nc", "localhost", "{}".format(local_port)], stdin=subprocess.PIPE,stdout=None) + try: + nc_send.communicate("Hello World".encode("utf-8"), timeout=1) + except subprocess.TimeoutExpired: + pass + nc_listen.terminate() + assert nc_listen.communicate()[0].decode("utf-8") == test_string + connection_localhost.remove_port_forward(port) + +def test_sshconnection_port_forward_remove_raise(connection_localhost): + port = 1337 + + with pytest.raises(ForwardError): + connection_localhost.remove_port_forward(port) + +def test_sshconnection_port_forward_add_duplicate(connection_localhost): + port = 1337 + + first_port = connection_localhost.add_port_forward(port) + second_port = connection_localhost.add_port_forward(port) + assert first_port == second_port + + +def test_sshconnection_put_file(connection_localhost, tmpdir): + port = 1337 + + p = tmpdir.join("config.yaml") + p.write( + """ +Teststring + """ + ) + connection_localhost.put_file(p, '/tmp/test') + +def test_sshconnection_get_file(connection_localhost, tmpdir): + + p = tmpdir.join("test") + connection_localhost.get_file('/tmp/test', p) + +def test_sshmanager_open(sshmanager): + con = sshmanager.open("localhost") + assert isinstance(con, SSHConnection) + +def test_sshmanager_add_forward(sshmanager): + port = sshmanager.request_forward("localhost", 3000) + assert port < 65536 + +def test_sshmanager_remove_forward(sshmanager): + sshmanager.remove_forward('localhost', 3000) + assert 3000 not in sshmanager.get('localhost')._forwards + +def test_sshmanager_close(sshmanager): + con = sshmanager.open("localhost") + assert isinstance(con, SSHConnection) + sshmanager.close("localhost") + +def test_sshmanager_remove_raise(sshmanager): + con = sshmanager.open("localhost") + con.connect() + with pytest.raises(ExecutionError): + sshmanager.remove_connection(con) + +def test_sshmanager_add_duplicate(sshmanager): + host = 'localhost' + con_there = sshmanager._connections[host] + con = SSHConnection(host) + sshmanager.add_connection(con) + con_now = sshmanager._connections[host] + assert con_now == con_there + +def test_sshmanager_add_new(sshmanager): + host = 'other_host' + con = SSHConnection(host) + sshmanager.add_connection(con) + con_now = sshmanager._connections[host] + assert con_now == con + + +def test_sshmanager_remove_raise(sshmanager): + con = SSHConnection("nosuchhost.notavailable") + with pytest.raises(ExecutionError): + con.connect() From 2c77c4557406f8d405ec83c4326200fd89bfce41 Mon Sep 17 00:00:00 2001 From: Rouven Czerwinski Date: Wed, 25 Apr 2018 17:20:26 +0200 Subject: [PATCH 4/4] doc/development: document SSHMANAGER Signed-off-by: Rouven Czerwinski --- doc/development.rst | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/doc/development.rst b/doc/development.rst index 0e7dc3719..e06f56de9 100644 --- a/doc/development.rst +++ b/doc/development.rst @@ -337,6 +337,34 @@ A root state is a state that has no dependencies. .. image:: res/graphstrategy-2.png +SSHManager +---------- + +Labgrid provides a SSHManager to allow connection reuse with control sockets. +To use the SSHManager in your code, import it from `labgrid.util.ssh`: + +.. code-block:: python + + from labgrid.util.ssh import SSHMANAGER + +you can now request or remove forwards: + +.. code-block:: python + + from labgrid.util.ssh import SSHMANAGER + + localport = SSHMANAGER.request_forward('somehost', 3000) + + SSHMANAGER.remove_forward('somehost', 3000) + +or get and put files:: + +.. code-block:: python + + from labgrid.util.ssh import SSHMANAGER + + SSHMANAGER.put_file('somehost', '/path/to/local/file', '/path/to/remote/file') + .. _contributing: Contributing