diff options
Diffstat (limited to 'tests/utils.py')
| -rw-r--r-- | tests/utils.py | 45 |
1 files changed, 45 insertions, 0 deletions
diff --git a/tests/utils.py b/tests/utils.py index c5d6c1c3..f57cd266 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -8,6 +8,7 @@ ############################################################################## import json import os +import paramiko def load_json_file(full_path): @@ -31,3 +32,47 @@ def write_json_file(full_path, data): with open(full_path, 'w+') as file: file.write(json.dumps(data)) + +class SSHClient(object): + def __init__(self, ip, username, password=None, pkey=None, + key_filename=None, log=None, look_for_keys=False, + allow_agent=False): + self.client = paramiko.SSHClient() + self.client.load_system_host_keys() + self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + self.client.connect(ip, username=username, password=password, + pkey=pkey, key_filename=key_filename, + look_for_keys=look_for_keys, + allow_agent=allow_agent) + self.log = log + + def __del__(self): + self.client.close() + + def ssh(self, command): + if self.log: + self.log.debug("Executing: %s" % command) + stdin, stdout, stderr = self.client.exec_command(command) + ret = stdout.channel.recv_exit_status() + output = list() + for line in stdout.read().splitlines(): + output.append(line) + if ret: + if self.log: + self.log.debug("*** FAILED to run command %s (%s)" % (command, ret)) + raise Exception( + "Unable to run \ncommand: %s\nret: %s" + % (command, ret)) + if self.log: + self.log.debug("*** SUCCESSFULLY run command %s" % command) + return ret, output + + def scp(self, source, dest, method='put'): + if self.log: + self.log.info("Copy %s -> %s" % (source, dest)) + ftp = self.client.open_sftp() + if method == 'put': + ftp.put(source, dest) + elif method == 'get': + ftp.get(source, dest) + ftp.close() |