-1

I need to create a method in python for remote SSH command execution that supports timeout (with partial output gathering from both stdout and stderr in case of timeout). I walked through lots of topics here and I couldn't find any full solution for this issue.

Below you'll find my suggestion on how to accomplish that in python 3.

Mirek
  • 185
  • 2
  • 11

1 Answers1

0

I'm using (slightly tuned) Thomas' solution for timeout, but another approach could be used if needed:

import paramiko
import signal

class Timeout:
    def __init__(self, seconds=1, error_message='Timeout', exception_type=TimeoutError):
        self.seconds = seconds
        self.error_message = error_message + " after {seconds}".format(seconds=seconds)
        self.exception_type = exception_type

    def handle_timeout(self, signum, frame):
        raise self.exception_type(self.error_message)

    def __enter__(self):
        signal.signal(signal.SIGALRM, self.handle_timeout)
        signal.alarm(self.seconds)

    def __exit__(self, type, value, traceback):
        signal.alarm(0)


def run_cmd(connection_info, cmd, timeout=5):
    """
    Runs command via SSH, supports timeout and partian stdout/stderr gathering.
    :param connection_info: Dict with 'hostname', 'username' and 'password' fields
    :param cmd: command to run
    :param timeout: timeout for command execution (does not apply to SSH session creation!)
    :return: stdout, stderr, exit code from command execution. Exit code is -1 for timed out commands
    """
    # Session creation can be done ous
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect(**connection_info)
    channel = client.get_transport().open_session()  # here TCP socket timeout could be tuned
    out = ""
    err = ""
    ext = -1  # exit code -1 for timed out commands
    channel.exec_command(cmd)
    try:
        with Timeout(timeout):
            while True:
                if channel.recv_ready():
                    out += channel.recv(1024).decode()
                if channel.recv_stderr_ready():
                    err += channel.recv_stderr(1024).decode()
                if channel.exit_status_ready() and not channel.recv_ready() and not channel.recv_stderr_ready():
                    ext = channel.recv_exit_status()
                    break
    except TimeoutError:
        print("command timeout")
    return out, err, ext

Now let's test this:

from time import time
def test(cmd, timeout):
    credentials = {"hostname": '127.0.0.1', "password": 'password', "username": 'user'}
    ts = time()
    print("\n---\nRunning: " + cmd)
    result = run_cmd(credentials, cmd, timeout)
    print("...it took: {0:4.2f}s".format(time() - ts))
    print("Result: {}\n".format(str(result)[:200]))

# Those should time out:
test('for i in {1..10}; do echo -n "OUT$i "; sleep 0.5; done', 2)
test('for i in {1..10}; do echo -n "ERR$i " >&2; sleep 0.5; done', 2)
test('for i in {1..10}; do echo -n "ERR$i " >&2; echo -n "OUT$i "; sleep 0.5; done', 2)

# Those should not time out:
test('for i in {1..10}; do echo -n "OUT$i "; sleep 0.1; done', 2)
test('for i in {1..10}; do echo -n "ERR$i " >&2; sleep 0.1; done', 2)
test('for i in {1..10}; do echo -n "ERR$i " >&2; echo -n "OUT$i "; sleep 0.1; done', 2)

# Large output testing, with timeout:
test("cat /dev/urandom | base64 |head -n 1000000", 2)
test("cat /dev/urandom | base64 |head -n 1000000 >&2", 2)
Mirek
  • 185
  • 2
  • 11