← Back to branch summary

~canonical-platform-qa/snappy-ecosystem-tests/remove-unneeded-FIXMEs

« back to all changes in this revision

Viewing changes to snappy_ecosystem_tests/utils/ssh.py

  • Committer: Omer Akram
  • Date: 2017-02-22 20:14:58 UTC
  • mfrom: (22.1.1 snappy-ecosystem-tests)
  • Revision ID: om26er@ubuntu.com-20170222201458-f0mfhw0ejl2w15bs
Merge with trunk

Show diffs side-by-side

added added

removed removed

Lines of Context:
snappy_ecosystem_tests/utils/ssh.py
 
1
# -*- Mode: Python; coding: utf-8; indent-tabs-mode: nil; tab-width: 4 -*-
 
2
 
 
3
#
 
4
# Snappy Ecosystem Tests
 
5
# Copyright (C) 2017 Canonical
 
6
#
 
7
# This program is free software: you can redistribute it and/or modify
 
8
# it under the terms of the GNU General Public License as published by
 
9
# the Free Software Foundation, either version 3 of the License, or
 
10
# (at your option) any later version.
 
11
#
 
12
# This program is distributed in the hope that it will be useful,
 
13
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
14
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
15
# GNU General Public License for more details.
 
16
#
 
17
# You should have received a copy of the GNU General Public License
 
18
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
19
#
 
20
 
 
21
"""Module to connect to ssh client for running tests."""
 
22
 
 
23
import paramiko
 
24
 
 
25
from snappy_ecosystem_tests.utils.user import get_remote_host_credentials
 
26
 
 
27
 
 
28
class SSHManager:
 
29
    """Manager class to keep a pool of ssh connections."""
 
30
    # Mangle the pool so that it cannot be easily touched from outside.
 
31
    __connection_pool = []
 
32
 
 
33
    def __init__(self):
 
34
        raise NotImplementedError(
 
35
            'Class cannot be instantiated, use get_instance() to get a '
 
36
            'SSHClient instance.')
 
37
 
 
38
    @staticmethod
 
39
    def get_instance(hostname, username, port=22, **kwargs):
 
40
        """Return the instance of SSHClient from a pool of active connections,
 
41
        otherwise create a new connection and return."""
 
42
        port = int(port)
 
43
        client = SSHManager.get_client(hostname, username, port)
 
44
        if client:
 
45
            if not client.get_transport().is_active():
 
46
                client.connect(
 
47
                    hostname, username=username, port=port, **kwargs)
 
48
        else:
 
49
            client = paramiko.SSHClient()
 
50
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
51
            client.connect(hostname, username=username, port=port, **kwargs)
 
52
            SSHManager.__connection_pool.append(client)
 
53
        return client
 
54
 
 
55
    @staticmethod
 
56
    def sync():
 
57
        """Update the connection pool to only keep alive connections."""
 
58
        pool_copy = SSHManager.__connection_pool.copy()
 
59
        for client in SSHManager.__connection_pool:
 
60
            if not client.get_transport().is_alive():
 
61
                pool_copy.remove(client)
 
62
        SSHManager.__connection_pool = pool_copy
 
63
 
 
64
    @staticmethod
 
65
    def get_client(hostname, username, port):
 
66
        """Return the matching client from the pool based on provided
 
67
        constraints."""
 
68
        connections = [
 
69
            client for client in SSHManager.__connection_pool if
 
70
            client.get_transport().getpeername() == (hostname, port) and
 
71
            username == client.get_transport().get_username()
 
72
        ]
 
73
        return connections[0] if connections else None
 
74
 
 
75
 
 
76
def run_command(command, hostname=None, username=None, port=None):
 
77
    """Run the given command on remote machine over ssh.
 
78
 
 
79
    :param command: a string of the command to run.
 
80
    :param hostname: The host to run command on.
 
81
    :param username: Name of the user on the remote host to login to.
 
82
    :param port: SSH port number.
 
83
    :raises ValueError: if command exits with non-zero status.
 
84
    :return: the stdout of the command.
 
85
    """
 
86
    _hostname, _username, _port = get_remote_host_credentials()
 
87
    ssh = SSHManager.get_instance(
 
88
        hostname or _hostname, username or _username, port or _port)
 
89
    _, stdout, stderr = ssh.exec_command(command)
 
90
    if stdout.channel.recv_exit_status() != 0:
 
91
        raise ValueError(stderr.read().decode())
 
92
    return stdout.read().decode()