ssh_helper.py 5.12 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22
#!/usr/bin/env python3

# Copyright (C) 2018 Simon Redman <sredman@cs.utah.edu>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import getpass
import sys
from pexpect import pxssh


23 24 25 26
DEFAULT_SSH_OPTIONS = {"StrictHostKeyChecking": "no",
                       "UserKnownHostsFile": "/dev/null"}


27
def log_in_password_failover(hostname, username, password="", ssh_options=DEFAULT_SSH_OPTIONS):
28 29 30 31 32 33
    """
    Attempt to log in a pxssh session using only ssh agent-provided public keys, then failover to requesting a password

    :param hostname: remote to connect to
    :param username: login username
    :param password: login password -- Will be prompted for if needed and not passed
34 35
    :param ssh_options: pxssh options, as would be in an ssh config file
    :return: connected pxssh session and the password used to log in to the host
36 37 38
    """
    # First, try in-memory authentication tools, including the password if it is defined
    try:
39 40
        session = pxssh.pxssh(options=ssh_options)
        session.force_password = False
41 42
        session.login(hostname, username, password)
    except pxssh.ExceptionPxssh:
43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86
        # That didn't work. Try a password.
        session = pxssh.pxssh(options=ssh_options)
        session.force_password = True # Probably never necessary -- we already tried key-based login
        password = getpass.getpass("Password for {user}@{host}: ".format(user=username, host=hostname))
        session.login(hostname, username, password)

    return session, password


def log_in_many_sessions(hostnames, usernames, passwords=None, ssh_options=DEFAULT_SSH_OPTIONS):
    """
    Attempt to log in a pxssh session for all requested hosts

    :param hostnames: list of hosts on which commands should be run
    :param usernames: list of usernames to use to log in to each host
    :param passwords: (optional) list of passwords to log in to each host
    :param ssh_options: options to pass to ssh, as per pxssh documentation
    :return: list of pxssh sessions and list of passwords which were used to log in to the hosts, where empty string means no password was used
    """
    sessions = []
    num_hosts = len(hostnames)
    assert len(usernames) == num_hosts, "Please provide one username for every host"
    if passwords is None:
        # Generate empty passwords for every host if no passwords were provided
        passwords = ["" for i in range(len(usernames))]
    assert len(passwords) == num_hosts, "If provided, there must be one password per username"

    for host_idx in range(0, num_hosts):
        hostname = hostnames[host_idx]
        username = usernames[host_idx]
        password = passwords[host_idx]

        if passwords == "":
            # To try to avoid prompting for passwords, try the previously-used password first
            password = passwords[host_idx - 1]
        try:
            session, returned_password = log_in_password_failover(hostname, username, password, ssh_options)
        except pxssh.ExceptionPxssh as e:
            print("Unable to log in to {user}@{host} using in-memory SSH keys nor provided password".format(user=username, host=hostname), file=sys.stderr)
            print(e, file=sys.stderr)
            continue

        sessions.append(session)
        passwords[host_idx] = returned_password
87

88
    return sessions, passwords
89 90


91
def get_output(session, encoding=sys.stdout.encoding):
92 93 94 95 96 97 98
    """
    Decode the raw bytes written by the SSH session

    :param session: pxssh session to read
    :param encoding: encoding to use to interpret the output
    :return: string-form of the read bytes
    """
99
    session.prompt()
100
    return str(session.before.decode(encoding))
101 102


103
def run_command_on_host(session, command):
104 105 106
    """
    Run a specified command on a single host

107
    :param session: logged-in pxssh session to run the command on
108
    :param command: command to run
109
    :return: None
110 111 112 113
    """
    session.sendline(command)


114
def run_commands_on_many_hosts(sessions, commands):
115 116 117
    """
    Add the specified command on each host in the network

118
    :param sessions: list of logged-in pxssh sessions to run commands on
119
    :param commands: list of commands to run, one per host
120
    :return: output from running each command, in the same order as the sessions were presented
121
    """
122 123 124
    outputs = []
    num_hosts = len(sessions)
    assert len(commands) == num_hosts, "Please provide one command for every session"
125 126

    for host_idx in range(0, num_hosts):
127
        session = sessions[host_idx]
128 129
        command = commands[host_idx]

130 131 132 133
        run_command_on_host(session, command)

    for host_idx in range(0, num_hosts):
        session = sessions[host_idx]
134

135 136
        output = get_output(session)
        outputs.append(output)
137

138
    return outputs