Commit fd74ec55 authored by Simon Redman's avatar Simon Redman

Add error-checking for ssh_helper

parent 3a68ae73
......@@ -18,7 +18,7 @@
import getpass
import sys
from pexpect import pxssh
from typing import List
from typing import List, Tuple, Optional
DEFAULT_SSH_OPTIONS = {"StrictHostKeyChecking": "no",
......@@ -141,17 +141,13 @@ def run_command_on_host(session, command):
session.sendline(command)
def run_commands_on_many_hosts(sessions, commands):
def _run_commands_on_many_hosts(sessions: List[pxssh.pxssh], commands: List[str]) -> List[str]:
"""
Add the specified command on each host in the network
:param sessions: list of logged-in pxssh sessions to run commands on
:param commands: list of commands to run, one per host
:return: output from running each command, in the same order as the sessions were presented
Run one command on each host in the network without error checking
"""
outputs = []
num_hosts = len(sessions)
assert len(commands) == num_hosts, "Please provide one command for every session"
assert len(commands) == len(sessions), "Please provide one command for every session"
for host_idx in range(0, num_hosts):
session = sessions[host_idx]
......@@ -168,13 +164,65 @@ def run_commands_on_many_hosts(sessions, commands):
return outputs
def run_commands_on_many_hosts(sessions: List[pxssh.pxssh], commands: List[str]) -> List[str]:
"""
Run the specified command on each host in the network
:param sessions: list of logged-in pxssh sessions to run commands on
:param commands: list of commands to run, one per host
:return: exit code and output from running each command, in the same order as the sessions were presented
"""
outputs = _run_commands_on_many_hosts(sessions, commands)
codes = get_exit_codes(sessions)
# Check for any error-indicating exit codes
errors = None
for index in range(0, len(codes)):
if codes[index] != 0:
errors = SSHCommandErrorError(session = sessions[index],
output = outputs[index],
code = codes[index],
next = errors)
if errors is not None: raise errors
return outputs
def get_exit_codes(sessions) -> List[int]:
"""
Get the exit code of the last command run in each session
"""
commands = ["echo $?" for session in sessions]
codes = run_commands_on_many_hosts(sessions, commands)
codes = [code.split()[2] for code in codes] # Get just the return code (not the echo'ed command)
codes = [int(code) for code in codes]
return codes
\ No newline at end of file
outputs = _run_commands_on_many_hosts(sessions, commands)
codes_strs = [code.split()[2] for code in outputs] # Get just the return code (not the echo'ed command)
codes = [int(code) for code in codes_strs]
return codes
class SSHError(Exception):
"""
Base class for errors related to SSH
"""
pass
class SSHCommandErrorError(SSHError):
"""
Raised to indicate one or more errors has occured while running a command
"""
def __init__(self, session: pxssh.pxssh, output: str, code:int, next: Optional):
"""
:param session: pxssh session which had the command run
:param output: output from the command which hopefully contains more error information
:param code: the exit code of the command
:param next: next exception, in case many commands were executed and more than one had a problem
"""
super(SSHCommandErrorError, self).__init__(output)
self.session = session
self.output = output
self.code = code
self.next = next
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment