Commit 66259e90 authored by Simon Redman's avatar Simon Redman

Spruce up ssh_helper's error checker

parent 797270a5
......@@ -165,7 +165,7 @@ def start_sniffer_on_network(graph: networkx.Graph,
# First the command is echoed
echos = ssh_helper.unchecked_run_commands_on_many_hosts(sessions, commands)
# Next we check for potential outputs
outputs = list(map(lambda s: ssh_helper.get_output(s, timeout=1), sessions))
outputs = list(map(lambda s: ssh_helper._get_output(s, timeout=1), sessions))
# Some real basic error checking
errors = None
......
......@@ -118,7 +118,7 @@ def network_graph_logout(netgraph):
del netgraph._node[node]['session']
def get_output(session, encoding=sys.stdout.encoding, timeout=None):
def _get_output(session, encoding=sys.stdout.encoding, timeout=None):
"""
Decode the raw bytes written by the SSH session
......@@ -131,9 +131,11 @@ def get_output(session, encoding=sys.stdout.encoding, timeout=None):
return str(session.before.decode(encoding))
def run_command_on_host(session, command):
def _run_command_on_host(session, command):
"""
Run a specified command on a single host
Note: be sure to match one command execution with one get_output, otherwise weird and
hard-to-understand things will happen
:param session: logged-in pxssh session to run the command on
:param command: command to run
......@@ -158,12 +160,12 @@ def unchecked_run_commands_on_many_hosts(sessions: List[pxssh.pxssh], commands:
session = sessions[host_idx]
command = commands[host_idx]
run_command_on_host(session, command)
_run_command_on_host(session, command)
for host_idx in range(0, num_hosts):
session = sessions[host_idx]
output = get_output(session)
output = _get_output(session)
outputs.append(output)
return outputs
......@@ -199,10 +201,9 @@ 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]
outputs = unchecked_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_strs = [list(filter(lambda line: len(line) > 0, code.split("\r\n")))[-1] for code in outputs] # Get just the return code (not the echo'ed command)
codes = [int(code) for code in codes_strs]
return codes
......
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