import errno

import logging

import paramiko

import socket

from os import path

from stat import S_ISDIR, S_ISREG

_KEY_TYPES = {

"dsa": paramiko.DSSKey,

"rsa": paramiko.RSAKey,

"ecdsa": paramiko.ECDSAKey,

"ed25519": paramiko.Ed25519Key,

}

class SFTPController(paramiko.SFTPClient):

def __init__(self, sock):

super().__init__(sock)

def exists(self, path):

try:

self.stat(path)

except IOError as e:

return e.errno != errno.ENOENT

return True

def list_dirs(self, path):

return [

d.filename for d in self.listdir_attr(path) if S_ISDIR(d.st_mode)

]

def list_files(self, path):

return [

f.filename for f in self.listdir_attr(path) if S_ISREG(f.st_mode)

]

@classmethod

def from_transport(cls, t):

chan = t.open_session()

chan.invoke_subsystem("sftp")

return cls(chan)

class SSHController:

def __init__(

self,

host,

user,

key_path=None,

key_password=None,

key_type="rsa",

ssh_password=None,

port=22,

):

self.host = host

self.user = user

self.ssh_password = ssh_password if key_path is None else None

self.port = port

self.nb_bytes = 1024

self.keys, self.transport = [], None

key_type = key_type.lower()

if key_path:

self.keys.append(

_KEY_TYPES[key_type].from_private_key(

open(path.expanduser(key_path), 'r'),

key_password,

)

)

elif ssh_password is None:

self.keys = paramiko.Agent().get_keys()

try:

key_file = _KEY_TYPES[key_type].from_private_key(

open(path.expanduser(f"~/.ssh/id_{key_type}"), 'r'),

key_password

)

except Exception:

pass

else:

self.keys.insert(

len(self.keys) if key_password is None else 0, key_file

)

if not self.keys:

logging.error("No valid key found")

def connect(self):

try:

ssh_socket = socket.create_connection((self.host, self.port))

except OSError as e:

logging.error(f"Connection failed: {e.strerror}")

return 1

self.transport = paramiko.Transport(ssh_socket)

if self.ssh_password is not None:

try:

self.transport.connect(

username=self.user,

password=self.ssh_password,

)

except paramiko.SSHException:

pass

else:

for key in self.keys:

try:

self.transport.connect(username=self.user, pkey=key)

except paramiko.SSHException:

continue

break

if not self.transport.is_authenticated():

logging.error("SSH negotiation failed")

return 1

logging.info(f"Successfully connected to {self.user}@{self.host}")

return 0

def __run_until_event(

self,

command,

stop_event,

display=True,

capture_output=False,

shell=True,

combine_stderr=False,

):

channel = self.transport.open_session()

output = ""

channel.settimeout(2)

channel.set_combine_stderr(combine_stderr)

if shell:

channel.get_pty()

channel.exec_command(command)

if not display and not capture_output:

stop_event.wait()

else:

while True:

try:

raw_data = channel.recv(self.nb_bytes)

except socket.timeout:

if stop_event.is_set():

break

continue

if not len(raw_data):

break

data = raw_data.decode("utf-8")

if display:

print(data, end='')

if capture_output:

output += data

if stop_event.is_set():

break

channel.close()

if not channel.exit_status_ready():

return (0, output.splitlines())

return (channel.recv_exit_status(), output.splitlines())

def __run_until_exit(

self,

command,

timeout,

display=True,

capture_output=False,

shell=True,

combine_stderr=False,

):

channel = self.transport.open_session()

output = ""

channel.settimeout(timeout)

channel.set_combine_stderr(combine_stderr)

if shell:

channel.get_pty()

channel.exec_command(command)

try:

if not display and not capture_output:

return (channel.recv_exit_status(), output.splitlines())

else:

while True:

raw_data = channel.recv(self.nb_bytes)

if not len(raw_data):

break

data = raw_data.decode("utf-8")

if display:

print(data, end='')

if capture_output:

output += data

except socket.timeout:

logging.warning(f"Timeout after {timeout}s")

return (1, output.splitlines())

except KeyboardInterrupt:

logging.info("KeyboardInterrupt")

return (0, output.splitlines())

finally:

channel.close()

return (channel.recv_exit_status(), output.splitlines())

def run(

self,

command,

display=False,

capture_output=False,

shell=True,

combine_stderr=False,

timeout=None,

stop_event=None,

):

if stop_event:

return self.__run_until_event(

command,

stop_event,

display=display,

shell=shell,

combine_stderr=combine_stderr,

capture_output=capture_output,

)

else:

return self.__run_until_exit(

command,

timeout,

display=display,

shell=shell,

combine_stderr=combine_stderr,

capture_output=capture_output,

)

def disconnect(self):

if self.transport:

self.transport.close()

def __getattr__(self, target):

def wrapper(*args, **kwargs):

if not self.transport.is_authenticated():

logging.error("SSH session is not ready")

return 1

sftp_channel = SFTPController.from_transport(self.transport)

r = getattr(sftp_channel, target)(*args, **kwargs)

sftp_channel.close()

return r

return wrapper