Skip to content
Snippets Groups Projects
Commit f115a6a7 authored by Mads M. Pedersen's avatar Mads M. Pedersen
Browse files

something

parent 8ef952ee
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
wetb/utils/cluster_tools/ssh_client.py
+ 376
− 358
View file @ f115a6a7
''' '''
Created on 27/11/2015 Created on 27/11/2015
@author: MMPE @author: MMPE
''' '''
from io import StringIO from io import StringIO
import paramiko import paramiko
import os import os
import sys import sys
import threading import threading
from _collections import deque from _collections import deque
import time import time
import traceback import traceback
import zipfile import zipfile
from wetb.utils.timing import print_time from wetb.utils.timing import print_time
import glob import glob
import getpass import getpass
from sshtunnel import SSHTunnelForwarder, SSH_CONFIG_FILE from sshtunnel import SSHTunnelForwarder, SSH_CONFIG_FILE
class SSHInteractiveAuthTunnelForwarder(SSHTunnelForwarder): class SSHInteractiveAuthTunnelForwarder(SSHTunnelForwarder):
def __init__( def __init__(
self, self,
interactive_auth_handler, interactive_auth_handler,
ssh_address_or_host=None, ssh_address_or_host=None,
ssh_config_file=SSH_CONFIG_FILE, ssh_config_file=SSH_CONFIG_FILE,
ssh_host_key=None, ssh_host_key=None,
ssh_password=None, ssh_password=None,
ssh_pkey=None, ssh_pkey=None,
ssh_private_key_password=None, ssh_private_key_password=None,
ssh_proxy=None, ssh_proxy=None,
ssh_proxy_enabled=True, ssh_proxy_enabled=True,
ssh_username=None, ssh_username=None,
local_bind_address=None, local_bind_address=None,
local_bind_addresses=None, local_bind_addresses=None,
logger=None, logger=None,
mute_exceptions=False, mute_exceptions=False,
remote_bind_address=None, remote_bind_address=None,
remote_bind_addresses=None, remote_bind_addresses=None,
set_keepalive=0.0, set_keepalive=0.0,
threaded=True, threaded=True,
compression=None, compression=None,
allow_agent=True, * allow_agent=True, *
args, ** args, **
kwargs): kwargs):
self.interactive_auth_handler = interactive_auth_handler self.interactive_auth_handler = interactive_auth_handler
SSHTunnelForwarder.__init__(self, ssh_address_or_host=ssh_address_or_host, ssh_config_file=ssh_config_file, ssh_host_key=ssh_host_key, ssh_password=ssh_password, ssh_pkey=ssh_pkey, ssh_private_key_password=ssh_private_key_password, ssh_proxy=ssh_proxy, ssh_proxy_enabled=ssh_proxy_enabled, ssh_username=ssh_username, local_bind_address=local_bind_address, local_bind_addresses=local_bind_addresses, logger=logger, mute_exceptions=mute_exceptions, remote_bind_address=remote_bind_address, remote_bind_addresses=remote_bind_addresses, set_keepalive=set_keepalive, threaded=threaded, compression=compression, allow_agent=allow_agent, *args, **kwargs) SSHTunnelForwarder.__init__(self, ssh_address_or_host=ssh_address_or_host, ssh_config_file=ssh_config_file, ssh_host_key=ssh_host_key, ssh_password=ssh_password, ssh_pkey=ssh_pkey, ssh_private_key_password=ssh_private_key_password, ssh_proxy=ssh_proxy, ssh_proxy_enabled=ssh_proxy_enabled, ssh_username=ssh_username, local_bind_address=local_bind_address, local_bind_addresses=local_bind_addresses, logger=logger, mute_exceptions=mute_exceptions, remote_bind_address=remote_bind_address, remote_bind_addresses=remote_bind_addresses, set_keepalive=set_keepalive, threaded=threaded, compression=compression, allow_agent=allow_agent, *args, **kwargs)
def _connect_to_gateway(self):
def _connect_to_gateway(self): """
""" Open connection to SSH gateway
Open connection to SSH gateway - First try with all keys loaded from an SSH agent (if allowed)
- First try with all keys loaded from an SSH agent (if allowed) - Then with those passed directly or read from ~/.ssh/config
- Then with those passed directly or read from ~/.ssh/config - As last resort, try with a provided password
- As last resort, try with a provided password """
""" try:
if self.ssh_password: # avoid conflict using both pass and pkey self._transport = self._get_transport()
self.logger.debug('Trying to log in with password: {0}' self._transport.start_client()
.format('*' * len(self.ssh_password))) self._transport.auth_interactive(self.ssh_username, self.interactive_auth_handler)
try: if self._transport.is_alive:
self._transport = self._get_transport() return
if self.interactive_auth_handler: except paramiko.AuthenticationException:
self._transport.start_client() self.logger.debug('Authentication error')
self._transport.auth_interactive(self.ssh_username, self.interactive_auth_handler) self._stop_transport()
else:
self._transport.connect(hostkey=self.ssh_host_key, self.logger.error('Could not open connection to gateway')
username=self.ssh_username,
password=self.ssh_password) def _connect_to_gateway_old(self):
"""
if self._transport.is_alive: Open connection to SSH gateway
return - First try with all keys loaded from an SSH agent (if allowed)
except paramiko.AuthenticationException: - Then with those passed directly or read from ~/.ssh/config
self.logger.debug('Authentication error') - As last resort, try with a provided password
self._stop_transport() """
if self.ssh_password: # avoid conflict using both pass and pkey
self.logger.debug('Trying to log in with password: {0}'
self.logger.error('Could not open connection to gateway') .format('*' * len(self.ssh_password)))
try:
class SSHClient(object): self._transport = self._get_transport()
"A wrapper of paramiko.SSHClient" if self.interactive_auth_handler:
TIMEOUT = 4 self._transport.start_client()
self._transport.auth_interactive(self.ssh_username, self.interactive_auth_handler)
def __init__(self, host, username, password=None, port=22, key=None, passphrase=None, interactive_auth_handler=None, gateway=None): else:
self.host = host self._transport.connect(hostkey=self.ssh_host_key,
self.username = username username=self.ssh_username,
self.password = password password=self.ssh_password)
self.port = port
self.key = key if self._transport.is_alive:
self.gateway=gateway return
self.interactive_auth_handler = interactive_auth_handler except paramiko.AuthenticationException:
self.disconnect = 0 self.logger.debug('Authentication error')
self.client = None self._stop_transport()
self.sftp = None
self.transport = None
if key is not None: self.logger.error('Could not open connection to gateway')
self.key = paramiko.RSAKey.from_private_key(StringIO(key), password=passphrase)
class SSHClient(object):
def info(self): "A wrapper of paramiko.SSHClient"
return self.host, self.username, self.password, self.port TIMEOUT = 4
def __enter__(self): def __init__(self, host, username, password=None, port=22, key=None, passphrase=None, interactive_auth_handler=None, gateway=None):
self.disconnect += 1 self.host = host
if self.client is None or self.client._transport is None or self.client._transport.is_active() is False: self.username = username
self.close() self.password = password
try: self.port = port
self.connect() self.key = key
self.disconnect = 1 self.gateway=gateway
except Exception as e: self.interactive_auth_handler = interactive_auth_handler
self.close() self.disconnect = 0
self.disconnect = 0 self.client = None
raise e self.sftp = None
return self.client self.transport = None
if key is not None:
def connect(self): self.key = paramiko.RSAKey.from_private_key(StringIO(key), password=passphrase)
# if self.password is None or self.password == "":
# raise IOError("Password not set for %s"%self.host) def info(self):
if self.gateway: return self.host, self.username, self.password, self.port
if self.gateway.interactive_auth_handler:
self.tunnel = SSHInteractiveAuthTunnelForwarder(self.gateway.interactive_auth_handler, def __enter__(self):
(self.gateway.host, self.gateway.port), self.disconnect += 1
ssh_username=self.gateway.username, if self.client is None or self.client._transport is None or self.client._transport.is_active() is False:
ssh_password=self.gateway.password, self.close()
remote_bind_address=(self.host, self.port), try:
local_bind_address=('0.0.0.0', 10022) self.connect()
) self.disconnect = 1
else: except Exception as e:
self.tunnel = SSHTunnelForwarder((self.gateway.host, self.gateway.port), self.close()
ssh_username=self.gateway.username, self.disconnect = 0
ssh_password=self.gateway.password, raise e
remote_bind_address=(self.host, self.port), return self.client
local_bind_address=('0.0.0.0', 10022)
) def connect(self):
# if self.password is None or self.password == "":
self.tunnel.start() # raise IOError("Password not set for %s"%self.host)
self.client = paramiko.SSHClient() if self.gateway:
self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) if self.gateway.interactive_auth_handler:
self.client.connect("127.0.0.1", 10022, username=self.username, password=self.password) self.tunnel = SSHInteractiveAuthTunnelForwarder(self.gateway.interactive_auth_handler,
(self.gateway.host, self.gateway.port),
ssh_username=self.gateway.username,
else: ssh_password=self.gateway.password,
self.client = paramiko.SSHClient() remote_bind_address=(self.host, self.port),
self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) local_bind_address=('0.0.0.0', 10022)
try: )
self.client.connect(self.host, self.port, username=self.username, password=self.password, pkey=self.key, timeout=self.TIMEOUT) else:
except paramiko.ssh_exception.SSHException as e: self.tunnel = SSHTunnelForwarder((self.gateway.host, self.gateway.port),
transport = self.client.get_transport() ssh_username=self.gateway.username,
transport.auth_interactive(self.username, self.interactive_auth_handler) ssh_password=self.gateway.password,
remote_bind_address=(self.host, self.port),
local_bind_address=('0.0.0.0', 10022)
)
assert self.client is not None
self.sftp = paramiko.SFTPClient.from_transport(self.client._transport) self.tunnel.start()
return self self.client = paramiko.SSHClient()
self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
def __exit__(self, *args): self.client.connect("127.0.0.1", 10022, username=self.username, password=self.password)
self.disconnect -= 1
if self.disconnect == 0:
self.close() else:
self.client = paramiko.SSHClient()
self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
def download(self, remotefilepath, localfile, verbose=False, retry=1): try:
if verbose: self.client.connect(self.host, self.port, username=self.username, password=self.password, pkey=self.key, timeout=self.TIMEOUT)
ret = None except paramiko.ssh_exception.SSHException as e:
print ("Download %s > %s" % (remotefilepath, str(localfile))) transport = self.client.get_transport()
with self: transport.auth_interactive(self.username, self.interactive_auth_handler)
for i in range(retry):
if i>0:
print ("Retry download %s, #%d"%(remotefilepath, i))
try: assert self.client is not None
if isinstance(localfile, (str, bytes, int)): self.sftp = paramiko.SFTPClient.from_transport(self.client._transport)
ret = self.sftp.get(remotefilepath, localfile) return self
elif hasattr(localfile, 'write'):
ret = self.sftp.putfo(remotefilepath, localfile) def __exit__(self, *args):
break self.disconnect -= 1
except: if self.disconnect == 0:
pass self.close()
print ("Download %s failed from %s"%(remotefilepath, self.host))
if verbose:
print (ret) def download(self, remotefilepath, localfile, verbose=False, retry=1):
if verbose:
def upload(self, localfile, filepath, verbose=False): ret = None
if verbose: print ("Download %s > %s" % (remotefilepath, str(localfile)))
print ("Upload %s > %s" % (localfile, filepath)) with self:
with self: for i in range(retry):
if isinstance(localfile, (str, bytes, int)): if i>0:
ret = self.sftp.put(localfile, filepath) print ("Retry download %s, #%d"%(remotefilepath, i))
elif hasattr(localfile, 'read'): try:
ret = self.sftp.putfo(localfile, filepath) if isinstance(localfile, (str, bytes, int)):
if verbose: ret = self.sftp.get(remotefilepath, localfile)
print (ret) elif hasattr(localfile, 'write'):
ret = self.sftp.putfo(remotefilepath, localfile)
break
def upload_files(self, localpath, remotepath, file_lst=["."], compression_level=1): except:
assert os.path.isdir(localpath) pass
if not isinstance(file_lst, (tuple, list)): print ("Download %s failed from %s"%(remotefilepath, self.host))
file_lst = [file_lst] if verbose:
files = ([os.path.join(root, f) for fp in file_lst for root,_,files in os.walk(os.path.join(localpath, fp )) for f in files] + print (ret)
[f for fp in file_lst for f in glob.glob(os.path.join(localpath, fp)) ])
files = set([os.path.abspath(f) for f in files]) def upload(self, localfile, filepath, verbose=False):
if verbose:
compression_levels = {0:zipfile.ZIP_STORED, 1:zipfile.ZIP_DEFLATED, 2:zipfile.ZIP_BZIP2, 3:zipfile.ZIP_LZMA} print ("Upload %s > %s" % (localfile, filepath))
zn = 'tmp_%s_%s.zip'%(id(self),time.time()) with self:
zipf = zipfile.ZipFile(zn, 'w', compression_levels[compression_level]) if isinstance(localfile, (str, bytes, int)):
try: ret = self.sftp.put(localfile, filepath)
for f in files: elif hasattr(localfile, 'read'):
zipf.write(f, os.path.relpath(f, localpath)) ret = self.sftp.putfo(localfile, filepath)
zipf.close() if verbose:
remote_zn = os.path.join(remotepath, zn).replace("\\","/") print (ret)
self.execute("mkdir -p %s"%(remotepath))
self.upload(zn, remote_zn) def upload_files(self, localpath, remotepath, file_lst=["."], compression_level=1):
self.execute("unzip %s -d %s && rm %s"%(remote_zn, remotepath, remote_zn)) assert os.path.isdir(localpath)
except: if not isinstance(file_lst, (tuple, list)):
raise file_lst = [file_lst]
finally: files = ([os.path.join(root, f) for fp in file_lst for root,_,files in os.walk(os.path.join(localpath, fp )) for f in files] +
os.remove(zn) [f for fp in file_lst for f in glob.glob(os.path.join(localpath, fp)) ])
files = set([os.path.abspath(f) for f in files])
def download_files(self, remote_path, localpath, file_lst=["."], compression_level=1): compression_levels = {0:zipfile.ZIP_STORED, 1:zipfile.ZIP_DEFLATED, 2:zipfile.ZIP_BZIP2, 3:zipfile.ZIP_LZMA}
if not isinstance(file_lst, (tuple, list)): zn = 'tmp_%s_%s.zip'%(id(self),time.time())
file_lst = [file_lst] zipf = zipfile.ZipFile(zn, 'w', compression_levels[compression_level])
file_lst = [f.replace("\\","/") for f in file_lst] try:
remote_zip = os.path.join(remote_path, "tmp.zip").replace("\\","/") for f in files:
self.execute("cd %s && zip -r tmp.zip %s"%(remote_path, " ".join(file_lst))) zipf.write(f, os.path.relpath(f, localpath))
zipf.close()
local_zip = os.path.join(localpath, "tmp.zip") remote_zn = os.path.join(remotepath, zn).replace("\\","/")
if not os.path.isdir(localpath): self.execute("mkdir -p %s"%(remotepath))
os.makedirs(localpath)
self.download(remote_zip, local_zip) self.upload(zn, remote_zn)
self.execute("rm -f %s" % remote_zip) self.execute("unzip %s -d %s && rm %s"%(remote_zn, remotepath, remote_zn))
with zipfile.ZipFile(local_zip, "r") as z: except:
z.extractall(localpath) raise
os.remove(local_zip) finally:
os.remove(zn)
def close(self):
for x in ["sftp", "client", 'tunnel' ]: def download_files(self, remote_path, localpath, file_lst=["."], compression_level=1):
try: if not isinstance(file_lst, (tuple, list)):
getattr(self, x).close() file_lst = [file_lst]
setattr(self, x, None) file_lst = [f.replace("\\","/") for f in file_lst]
except: remote_zip = os.path.join(remote_path, "tmp.zip").replace("\\","/")
pass self.execute("cd %s && zip -r tmp.zip %s"%(remote_path, " ".join(file_lst)))
self.disconnect = False
local_zip = os.path.join(localpath, "tmp.zip")
def file_exists(self, filename): if not os.path.isdir(localpath):
_, out, _ = (self.execute('[ -f %s ] && echo "File exists" || echo "File does not exists"' % filename.replace("\\", "/"))) os.makedirs(localpath)
return out.strip() == "File exists" self.download(remote_zip, local_zip)
self.execute("rm -f %s" % remote_zip)
def execute(self, command, sudo=False, verbose=False): with zipfile.ZipFile(local_zip, "r") as z:
feed_password = False z.extractall(localpath)
if sudo and self.username != "root": os.remove(local_zip)
command = "sudo -S -p '' %s" % command
feed_password = self.password is not None and len(self.password) > 0
if isinstance(command, (list, tuple)): def close(self):
command = "\n".join(command) for x in ["sftp", "client", 'tunnel' ]:
try:
if verbose: getattr(self, x).close()
print (">>> " + command) setattr(self, x, None)
with self as ssh: except:
if ssh is None: pass
exc_info = sys.exc_info() self.disconnect = False
traceback.print_exception(*exc_info)
raise Exception("ssh_client exe ssh is None") def file_exists(self, filename):
stdin, stdout, stderr = ssh.exec_command(command) _, out, _ = (self.execute('[ -f %s ] && echo "File exists" || echo "File does not exists"' % filename.replace("\\", "/")))
if feed_password: return out.strip() == "File exists"
stdin.write(self.password + "\n")
stdin.flush() def execute(self, command, sudo=False, verbose=False):
feed_password = False
v, out, err = stdout.channel.recv_exit_status(), stdout.read().decode(), stderr.read().decode() if sudo and self.username != "root":
command = "sudo -S -p '' %s" % command
feed_password = self.password is not None and len(self.password) > 0
if v: if isinstance(command, (list, tuple)):
raise Warning ("out:\n%s\n----------\nerr:\n%s" % (out, err)) command = "\n".join(command)
elif verbose:
if out: if verbose:
sys.stdout.write(out) print (">>> " + command)
if err: with self as ssh:
sys.stderr.write(err) if ssh is None:
return v, out, err exc_info = sys.exc_info()
traceback.print_exception(*exc_info)
def append_wine_path(self, path): raise Exception("ssh_client exe ssh is None")
ret = self.execute('wine regedit /E tmp.reg "HKEY_LOCAL_MACHINE\System\CurrentControlSet\Control\Session Manager\Environment"') stdin, stdout, stderr = ssh.exec_command(command)
self.download('tmp.reg', 'tmp.reg') if feed_password:
with open('tmp.reg') as fid: stdin.write(self.password + "\n")
lines = fid.readlines() stdin.flush()
path_line = [l for l in lines if l.startswith('"PATH"=')][0] v, out, err = stdout.channel.recv_exit_status(), stdout.read().decode(), stderr.read().decode()
for p in path_line[8:-1].split(";"):
if os.path.abspath(p) == os.path.abspath(p):
return if v:
if path not in path_line: raise Warning ("out:\n%s\n----------\nerr:\n%s" % (out, err))
path_line = path_line.strip()[:-1] + ";" + path + '"' elif verbose:
if out:
with open('tmp.reg', 'w') as fid: sys.stdout.write(out)
fid.write("".join(lines[:3] + [path_line])) if err:
self.upload('tmp.reg', 'tmp.reg') sys.stderr.write(err)
ret = self.execute('wine regedit tmp.reg') return v, out, err
def glob(self, filepattern, cwd="", recursive=False): def append_wine_path(self, path):
cwd = os.path.join(cwd, os.path.split(filepattern)[0]).replace("\\", "/") ret = self.execute('wine regedit /E tmp.reg "HKEY_LOCAL_MACHINE\System\CurrentControlSet\Control\Session Manager\Environment"')
filepattern = os.path.split(filepattern)[1] self.download('tmp.reg', 'tmp.reg')
if recursive: with open('tmp.reg') as fid:
_, out, _ = self.execute(r'find %s -type f -name "%s"' % (cwd, filepattern)) lines = fid.readlines()
else:
_, out, _ = self.execute(r'find %s -maxdepth 1 -type f -name "%s"' % (cwd, filepattern)) path_line = [l for l in lines if l.startswith('"PATH"=')][0]
return [file for file in out.strip().split("\n") if file != ""] for p in path_line[8:-1].split(";"):
if os.path.abspath(p) == os.path.abspath(p):
return
if path not in path_line:
path_line = path_line.strip()[:-1] + ";" + path + '"'
class SharedSSHClient(SSHClient):
def __init__(self, host, username, password=None, port=22, key=None, passphrase=None, interactive_auth_handler=None, gateway=None): with open('tmp.reg', 'w') as fid:
SSHClient.__init__(self, host, username, password=password, port=port, key=key, passphrase=passphrase, interactive_auth_handler=interactive_auth_handler, gateway=gateway) fid.write("".join(lines[:3] + [path_line]))
self.shared_ssh_lock = threading.RLock() self.upload('tmp.reg', 'tmp.reg')
self.shared_ssh_queue = deque() ret = self.execute('wine regedit tmp.reg')
self.next = None
def glob(self, filepattern, cwd="", recursive=False):
cwd = os.path.join(cwd, os.path.split(filepattern)[0]).replace("\\", "/")
def execute(self, command, sudo=False, verbose=False): filepattern = os.path.split(filepattern)[1]
res = SSHClient.execute(self, command, sudo=sudo, verbose=verbose) if recursive:
return res _, out, _ = self.execute(r'find %s -type f -name "%s"' % (cwd, filepattern))
else:
def __enter__(self): _, out, _ = self.execute(r'find %s -maxdepth 1 -type f -name "%s"' % (cwd, filepattern))
with self.shared_ssh_lock: return [file for file in out.strip().split("\n") if file != ""]
if self.next == threading.currentThread():
return self.client
self.shared_ssh_queue.append(threading.current_thread())
if self.next is None:
self.next = self.shared_ssh_queue.popleft() class SharedSSHClient(SSHClient):
def __init__(self, host, username, password=None, port=22, key=None, passphrase=None, interactive_auth_handler=None, gateway=None):
while self.next != threading.currentThread(): SSHClient.__init__(self, host, username, password=password, port=port, key=key, passphrase=passphrase, interactive_auth_handler=interactive_auth_handler, gateway=gateway)
time.sleep(1) self.shared_ssh_lock = threading.RLock()
SSHClient.__enter__(self) self.shared_ssh_queue = deque()
return self.client self.next = None
def __exit__(self, *args):
with self.shared_ssh_lock: def execute(self, command, sudo=False, verbose=False):
if next != threading.current_thread(): res = SSHClient.execute(self, command, sudo=sudo, verbose=verbose)
with self.shared_ssh_lock: return res
if len(self.shared_ssh_queue) > 0:
self.next = self.shared_ssh_queue.popleft() def __enter__(self):
else: with self.shared_ssh_lock:
self.next = None if self.next == threading.currentThread():
return self.client
self.shared_ssh_queue.append(threading.current_thread())
if self.next is None:
if __name__ == "__main__": self.next = self.shared_ssh_queue.popleft()
from mmpe.ui.qt_ui import QtInputUI
q = QtInputUI(None) while self.next != threading.currentThread():
x = None time.sleep(1)
username, password = "mmpe", x.password #q.get_login("mmpe") SSHClient.__enter__(self)
return self.client
client = SSHClient(host='gorm', port=22, username=username, password=password) def __exit__(self, *args):
print (client.glob("*.*", ".hawc2launcher/medium1__1__")) with self.shared_ssh_lock:
# ssh.upload('../News.txt', 'news.txt') if next != threading.current_thread():
with self.shared_ssh_lock:
if len(self.shared_ssh_queue) > 0:
self.next = self.shared_ssh_queue.popleft()
else:
self.next = None
if __name__ == "__main__":
from mmpe.ui.qt_ui import QtInputUI
q = QtInputUI(None)
x = None
username, password = "mmpe", x.password #q.get_login("mmpe")
client = SSHClient(host='gorm', port=22, username=username, password=password)
print (client.glob("*.*", ".hawc2launcher/medium1__1__"))
# ssh.upload('../News.txt', 'news.txt')
wetb/utils/tests/test_ssh_client.py
+ 196
− 196
View file @ f115a6a7
''' '''
Created on 23. dec. 2016 Created on 23. dec. 2016
@author: mmpe @author: mmpe
''' '''
import unittest import unittest
from wetb.utils.cluster_tools.ssh_client import SSHClient from wetb.utils.cluster_tools.ssh_client import SSHClient
import sys import sys
import os import os
try: try:
import x import x
except: except:
x=None x=None
import io import io
from wetb.utils.timing import print_time from wetb.utils.timing import print_time
import shutil import shutil
import paramiko import paramiko
from paramiko.message import Message from paramiko.message import Message
from paramiko.common import cMSG_SERVICE_REQUEST from paramiko.common import cMSG_SERVICE_REQUEST
import getpass import getpass
import logging import logging
import getpass import getpass
class sshrisoe_interactive_auth_handler(object): class sshrisoe_interactive_auth_handler(object):
def __init__(self, password): def __init__(self, password):
self.password = password self.password = password
def __call__(self, title, instructions, prompt_list): def __call__(self, title, instructions, prompt_list):
if prompt_list: if prompt_list:
if prompt_list[0][0]=="AD Password: ": if prompt_list[0][0]=="AD Password: ":
return [self.password] return [self.password]
return [getpass.getpass(prompt_list[0][0])] return [getpass.getpass(prompt_list[0][0])]
return [] return []
tfp = os.path.join(os.path.dirname(__file__), 'test_files/') tfp = os.path.join(os.path.dirname(__file__), 'test_files/')
all = 0 all = 0
class TestSSHClient(unittest.TestCase): class TestSSHClient(unittest.TestCase):
def setUp(self): def setUp(self):
if x: if x:
self.ssh = SSHClient('gorm', 'mmpe',x.mmpe ) self.ssh = SSHClient('gorm', 'mmpe',x.mmpe )
def test_execute(self): def test_execute(self):
if 0 or all: if 0 or all:
if x: if x:
_,out,_ = self.ssh.execute("ls -a") _,out,_ = self.ssh.execute("ls -a")
ssh_ls = ";".join(sorted(out.split("\n"))[3:]) #Exclude ['', '.', '..'] ssh_ls = ";".join(sorted(out.split("\n"))[3:]) #Exclude ['', '.', '..']
win_ls = ";".join(sorted(os.listdir(r"z:"))) win_ls = ";".join(sorted(os.listdir(r"z:")))
self.assertEqual(ssh_ls, win_ls) self.assertEqual(ssh_ls, win_ls)
def test_file_transfer(self): def test_file_transfer(self):
if 0 or all: if 0 or all:
if x: if x:
self.ssh.execute("rm -f tmp.txt") self.ssh.execute("rm -f tmp.txt")
io.StringIO() io.StringIO()
txt = "Hello world" txt = "Hello world"
f = io.StringIO(txt) f = io.StringIO(txt)
f.seek(0) f.seek(0)
self.ssh.upload(f, "tmp.txt") self.ssh.upload(f, "tmp.txt")
_,out,_ = self.ssh.execute("cat tmp.txt") _,out,_ = self.ssh.execute("cat tmp.txt")
self.assertEqual(out, txt) self.assertEqual(out, txt)
fn = tfp + "tmp.txt" fn = tfp + "tmp.txt"
if os.path.isfile (fn): if os.path.isfile (fn):
os.remove(fn) os.remove(fn)
self.assertFalse(os.path.isfile(fn)) self.assertFalse(os.path.isfile(fn))
self.ssh.download("tmp.txt", fn) self.ssh.download("tmp.txt", fn)
with open(fn) as fid: with open(fn) as fid:
self.assertEqual(fid.read(), txt) self.assertEqual(fid.read(), txt)
def test_folder_transfer(self): def test_folder_transfer(self):
if 0 or all: if 0 or all:
if x: if x:
p = r"C:\mmpe\HAWC2\models\version_12.3beta/" p = r"C:\mmpe\HAWC2\models\version_12.3beta/"
p = r'C:\mmpe\programming\python\WindEnergyToolbox\wetb\hawc2\tests\test_files\simulation_setup\DTU10MWRef6.0_IOS/' p = r'C:\mmpe\programming\python\WindEnergyToolbox\wetb\hawc2\tests\test_files\simulation_setup\DTU10MWRef6.0_IOS/'
self.ssh.execute("rm -r -f ./tmp_test") self.ssh.execute("rm -r -f ./tmp_test")
self.ssh.upload_files(p, "./tmp_test", ["input/"]) self.ssh.upload_files(p, "./tmp_test", ["input/"])
shutil.rmtree("./test/input", ignore_errors=True) shutil.rmtree("./test/input", ignore_errors=True)
self.ssh.download_files("./tmp_test", tfp, "input/" ) self.ssh.download_files("./tmp_test", tfp, "input/" )
os.path.isfile(tfp + "/input/data/DTU_10MW_RWT_Blade_st.dat") os.path.isfile(tfp + "/input/data/DTU_10MW_RWT_Blade_st.dat")
shutil.rmtree("./test/input", ignore_errors=True) shutil.rmtree("./test/input", ignore_errors=True)
def test_folder_transfer_specific_files_uppercase(self): def test_folder_transfer_specific_files_uppercase(self):
if 0 or all: if 0 or all:
if x: if x:
p = tfp p = tfp
files = [os.path.join(tfp, "TEST.txt")] files = [os.path.join(tfp, "TEST.txt")]
self.ssh.execute("rm -r -f ./tmp_test") self.ssh.execute("rm -r -f ./tmp_test")
self.ssh.upload_files(p, "./tmp_test", file_lst=files) self.ssh.upload_files(p, "./tmp_test", file_lst=files)
self.assertFalse(self.ssh.file_exists("./tmp_test/test.txt")) self.assertFalse(self.ssh.file_exists("./tmp_test/test.txt"))
self.assertTrue(self.ssh.file_exists("./tmp_test/TEST.txt")) self.assertTrue(self.ssh.file_exists("./tmp_test/TEST.txt"))
def test_folder_transfer_specific_files(self): def test_folder_transfer_specific_files(self):
if 0 or all: if 0 or all:
if x: if x:
p = r"C:\mmpe\HAWC2\models\version_12.3beta/" p = r"C:\mmpe\HAWC2\models\version_12.3beta/"
p = r'C:\mmpe\programming\python\WindEnergyToolbox\wetb\hawc2\tests\test_files\simulation_setup\DTU10MWRef6.0_IOS/' p = r'C:\mmpe\programming\python\WindEnergyToolbox\wetb\hawc2\tests\test_files\simulation_setup\DTU10MWRef6.0_IOS/'
files = [os.path.join(os.path.relpath(root, p), f) for root,_,files in os.walk(p+"input/") for f in files] files = [os.path.join(os.path.relpath(root, p), f) for root,_,files in os.walk(p+"input/") for f in files]
self.ssh.execute("rm -r -f ./tmp_test") self.ssh.execute("rm -r -f ./tmp_test")
self.ssh.upload_files(p, "./tmp_test", file_lst=files[:5]) self.ssh.upload_files(p, "./tmp_test", file_lst=files[:5])
self.ssh.download_files("./tmp_test", tfp + "tmp/", file_lst = files[:3]) self.ssh.download_files("./tmp_test", tfp + "tmp/", file_lst = files[:3])
self.assertEqual(len(os.listdir(tfp+"tmp/input/data/")),2) self.assertEqual(len(os.listdir(tfp+"tmp/input/data/")),2)
shutil.rmtree(tfp + "tmp/") shutil.rmtree(tfp + "tmp/")
# def test_ssh_gorm(self): # def test_ssh_gorm(self):
# if x: # if x:
# ssh = SSHClient('gorm.risoe.dk', 'mmpe', x.mmpe) # ssh = SSHClient('gorm.risoe.dk', 'mmpe', x.mmpe)
# _,out,_ = ssh.execute("hostname") # _,out,_ = ssh.execute("hostname")
# self.assertEqual(out.strip(), "g-000.risoe.dk") # self.assertEqual(out.strip(), "g-000.risoe.dk")
# def test_ssh_g047(self): # def test_ssh_g047(self):
# if x: # if x:
# gateway = SSHClient('gorm.risoe.dk', 'mmpe', x.mmpe) # gateway = SSHClient('gorm.risoe.dk', 'mmpe', x.mmpe)
# ssh = SSHClient('g-047', "mmpe", x.mmpe, gateway=gateway) # ssh = SSHClient('g-047', "mmpe", x.mmpe, gateway=gateway)
# self.assertEqual(ssh.execute('hostname')[1].strip(), "g-047") # self.assertEqual(ssh.execute('hostname')[1].strip(), "g-047")
def test_ssh_risoe(self): # def test_ssh_risoe(self):
if x: # if x:
#
ssh = SSHClient('ssh.risoe.dk', 'mmpe', interactive_auth_handler = sshrisoe_interactive_auth_handler(x.mmpe)) # ssh = SSHClient('ssh.risoe.dk', 'mmpe', interactive_auth_handler = sshrisoe_interactive_auth_handler(x.mmpe))
_,out,_ = ssh.execute("hostname") # _,out,_ = ssh.execute("hostname")
self.assertEqual(out.strip(), "ssh-03.risoe.dk") # self.assertEqual(out.strip(), "ssh-03.risoe.dk")
def test_ssh_risoe_gorm(self): def test_ssh_risoe_gorm(self):
if x: if x:
gateway = SSHClient('ssh.risoe.dk', 'mmpe', interactive_auth_handler = sshrisoe_interactive_auth_handler(x.mmpe)) gateway = SSHClient('ssh.risoe.dk', 'mmpe', interactive_auth_handler = sshrisoe_interactive_auth_handler(x.mmpe))
ssh = SSHClient('gorm.risoe.dk', 'mmpe', x.mmpe, gateway = gateway) ssh = SSHClient('gorm.risoe.dk', 'mmpe', x.mmpe, gateway = gateway)
_,out,_ = ssh.execute("hostname") _,out,_ = ssh.execute("hostname")
self.assertEqual(out.strip(), "g-000.risoe.dk") self.assertEqual(out.strip(), "g-000.risoe.dk")
# def test_ssh_risoe(self): # def test_ssh_risoe(self):
# #logger = logging.getLogger("paramiko") # #logger = logging.getLogger("paramiko")
# #logger.setLevel(logging.DEBUG) # for example # #logger.setLevel(logging.DEBUG) # for example
# #ch = logging.StreamHandler(sys.stdout) # #ch = logging.StreamHandler(sys.stdout)
# #ch.setLevel(logging.DEBUG) # #ch.setLevel(logging.DEBUG)
# #formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') # #formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
# #ch.setFormatter(formatter) # #ch.setFormatter(formatter)
# #logger.addHandler(ch) # #logger.addHandler(ch)
# ssh = SSHClient('ssh.risoe.dk', 'mmpe') # ssh = SSHClient('ssh.risoe.dk', 'mmpe')
# print (ssh.connect()) # print (ssh.connect())
# return # return
# username = "mmpe" # username = "mmpe"
# #
# client = paramiko.client.SSHClient() # client = paramiko.client.SSHClient()
# #
# # Any means of getting the PKey will do. This code assumes you've only got one key loaded in your active ssh-agent. # # Any means of getting the PKey will do. This code assumes you've only got one key loaded in your active ssh-agent.
# # See also: # # See also:
# # - http://docs.paramiko.org/en/1.17/api/keys.html#paramiko.pkey.PKey # # - http://docs.paramiko.org/en/1.17/api/keys.html#paramiko.pkey.PKey
# # - http://docs.paramiko.org/en/1.17/api/client.html#paramiko.client.SSHClient.connect # # - http://docs.paramiko.org/en/1.17/api/client.html#paramiko.client.SSHClient.connect
# my_pkey = None #paramiko.agent.Agent().get_keys()[0] # my_pkey = None #paramiko.agent.Agent().get_keys()[0]
# #
# try: # try:
# client.connect( # client.connect(
# hostname="ssh.risoe.dk", # hostname="ssh.risoe.dk",
# port=22, # port=22,
# username=username, # username=username,
# look_for_keys=False, # look_for_keys=False,
# pkey=my_pkey # pkey=my_pkey
# ) # )
# except paramiko.ssh_exception.SSHException as e: # except paramiko.ssh_exception.SSHException as e:
# pass # pass
# #
# transport = client.get_transport() # transport = client.get_transport()
# #
# # Sometimes sshd is configured to use 'keyboard-interactive' instead of 'password' to implement the YubiKey challenge. # # Sometimes sshd is configured to use 'keyboard-interactive' instead of 'password' to implement the YubiKey challenge.
# # In that case, you can use something like this. # # In that case, you can use something like this.
# # The code below assumes the server will only ask one question and expect the YubiKey OTP as an answer. # # The code below assumes the server will only ask one question and expect the YubiKey OTP as an answer.
# # If there's more questions to answer, you should handle those per the docs at: # # If there's more questions to answer, you should handle those per the docs at:
# # http://docs.paramiko.org/en/1.17/api/transport.html#paramiko.transport.Transport.auth_interactive # # http://docs.paramiko.org/en/1.17/api/transport.html#paramiko.transport.Transport.auth_interactive
# # # #
# def interactive_handler(title, instructions, prompt_list): # def interactive_handler(title, instructions, prompt_list):
# if prompt_list: # if prompt_list:
# if prompt_list[0][0]=="AD Password: ": # if prompt_list[0][0]=="AD Password: ":
# return [x.mmpe] # return [x.mmpe]
# return [getpass.getpass(prompt_list[0][0])] # return [getpass.getpass(prompt_list[0][0])]
# print ("here") # print ("here")
# return [] # return []
# transport.auth_interactive(username, interactive_handler) # transport.auth_interactive(username, interactive_handler)
# #
# #transport.auth_password(username, x.mmpe) # #transport.auth_password(username, x.mmpe)
# #
# # You should now be able to use client as the authenticated user. # # You should now be able to use client as the authenticated user.
# client.exec_command("echo hej") # client.exec_command("echo hej")
# #
if __name__ == "__main__": if __name__ == "__main__":
#import sys;sys.argv = ['', 'Test.testName'] #import sys;sys.argv = ['', 'Test.testName']
unittest.main() unittest.main()
\ No newline at end of file
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