Refactor sftp plugin, to avoid duplicate code

This commit is contained in:
Mathias Petermann 2020-11-11 11:58:36 +01:00
parent 648565d3e9
commit 2b01bdb471

View File

@ -19,7 +19,7 @@
from os import path
import paramiko
from paramiko.ssh_exception import AuthenticationException
from paramiko.ssh_exception import AuthenticationException, NoValidConnectionsError
from flask import current_app
from lemur.plugins import lemur_sftp
@ -97,6 +97,50 @@ class SFTPDestinationPlugin(DestinationPlugin):
},
]
def open_sftp_connection(self, options):
host = self.get_option("host", options)
port = self.get_option("port", options)
user = self.get_option("user", options)
password = self.get_option("password", options)
ssh_priv_key = self.get_option("privateKeyPath", options)
ssh_priv_key_pass = self.get_option("privateKeyPass", options)
# delete files
try:
current_app.logger.debug(
"Connecting to {0}@{1}:{2}".format(user, host, port)
)
ssh = paramiko.SSHClient()
# allow connection to the new unknown host
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
# open the ssh connection
if password:
current_app.logger.debug("Using password")
ssh.connect(host, username=user, port=port, password=password)
elif ssh_priv_key:
current_app.logger.debug("Using RSA private key")
pkey = paramiko.RSAKey.from_private_key_file(
ssh_priv_key, ssh_priv_key_pass
)
ssh.connect(host, username=user, port=port, pkey=pkey)
else:
current_app.logger.error(
"No password or private key provided. Can't proceed"
)
raise AuthenticationException
# open the sftp session inside the ssh connection
return ssh.open_sftp(), ssh
except AuthenticationException as e:
current_app.logger.error("ERROR in {0}: {1}".format(e.__class__, e))
raise AuthenticationException("Couldn't connect to {0}, due to an Authentication exception.")
except NoValidConnectionsError as e:
current_app.logger.error("ERROR in {0}: {1}".format(e.__class__, e))
raise NoValidConnectionsError("Couldn't connect to {0}, possible timeout or invalid hostname")
# this is called when using this as a default destination plugin
def upload(self, name, body, private_key, cert_chain, options, **kwargs):
@ -149,41 +193,9 @@ class SFTPDestinationPlugin(DestinationPlugin):
# here the file is deleted
def delete_file(self, dst_path, files, options):
host = self.get_option("host", options)
port = self.get_option("port", options)
user = self.get_option("user", options)
password = self.get_option("password", options)
ssh_priv_key = self.get_option("privateKeyPath", options)
ssh_priv_key_pass = self.get_option("privateKeyPass", options)
# delete files
try:
current_app.logger.debug(
"Connecting to {0}@{1}:{2}".format(user, host, port)
)
ssh = paramiko.SSHClient()
# allow connection to the new unknown host
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
# open the ssh connection
if password:
current_app.logger.debug("Using password")
ssh.connect(host, username=user, port=port, password=password)
elif ssh_priv_key:
current_app.logger.debug("Using RSA private key")
pkey = paramiko.RSAKey.from_private_key_file(
ssh_priv_key, ssh_priv_key_pass
)
ssh.connect(host, username=user, port=port, pkey=pkey)
else:
current_app.logger.error(
"No password or private key provided. Can't proceed"
)
raise AuthenticationException
# open the sftp session inside the ssh connection
sftp = ssh.open_sftp()
# open the ssh and sftp sessions
sftp, ssh = self.open_sftp_connection(options)
# delete files
for filename, _ in files.items():
@ -202,7 +214,8 @@ class SFTPDestinationPlugin(DestinationPlugin):
sftp.remove(path.join(dst_path, filename))
ssh.close()
except (AuthenticationException, NoValidConnectionsError) as e:
raise e
except Exception as e:
current_app.logger.error("ERROR in {0}: {1}".format(e.__class__, e))
try:
@ -213,38 +226,9 @@ class SFTPDestinationPlugin(DestinationPlugin):
# here the file is uploaded for real, this helps to keep this class DRY
def upload_file(self, dst_path, files, options):
host = self.get_option("host", options)
port = self.get_option("port", options)
user = self.get_option("user", options)
password = self.get_option("password", options)
ssh_priv_key = self.get_option("privateKeyPath", options)
ssh_priv_key_pass = self.get_option("privateKeyPass", options)
# upload files
try:
current_app.logger.debug(
"Connecting to {0}@{1}:{2}".format(user, host, port)
)
ssh = paramiko.SSHClient()
# allow connection to the new unknown host
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
# open the ssh connection
if password:
current_app.logger.debug("Using password")
ssh.connect(host, username=user, port=port, password=password)
elif ssh_priv_key:
current_app.logger.debug("Using RSA private key")
pkey = paramiko.RSAKey.from_private_key_file(
ssh_priv_key, ssh_priv_key_pass
)
ssh.connect(host, username=user, port=port, pkey=pkey)
else:
current_app.logger.error(
"No password or private key provided. Can't proceed"
)
raise AuthenticationException
# open the ssh and sftp sessions
sftp, ssh = self.open_sftp_connection(options)
# split the path into it's segments, so we can create it recursively
allparts = []
@ -261,9 +245,6 @@ class SFTPDestinationPlugin(DestinationPlugin):
path_copy = parts[0]
allparts.insert(0, parts[1])
# open the sftp session inside the ssh connection
sftp = ssh.open_sftp()
# make sure that the destination path exists, recursively
remote_path = allparts[0]
for part in allparts:
@ -301,9 +282,8 @@ class SFTPDestinationPlugin(DestinationPlugin):
ssh.close()
except AuthenticationException as e:
current_app.logger.error("ERROR in {0}: {1}".format(e.__class__, e))
raise AuthenticationException("Couldn't connect to {0}, due to an Authentication exception.")
except (AuthenticationException, NoValidConnectionsError) as e:
raise e
except Exception as e:
current_app.logger.error("ERROR in {0}: {1}".format(e.__class__, e))
try:
@ -314,4 +294,5 @@ class SFTPDestinationPlugin(DestinationPlugin):
if hasattr(e, 'errors'):
for _, error in e.errors.items():
message = error.strerror
raise Exception('Couldn\'t upload file to {}, error message: {}'.format(host, message))
raise Exception(
'Couldn\'t upload file to {}, error message: {}'.format(self.get_option("host", options), message))