Add SFTP client
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
ci/woodpecker/tag/woodpecker Pipeline was successful

This commit is contained in:
Benjamin Renard 2022-06-28 11:05:43 +02:00
parent b80cc3b3b6
commit 9511b31a79
4 changed files with 266 additions and 2 deletions

View file

@ -7,6 +7,7 @@ import getpass
import logging
import socket
import sys
import os.path
log = logging.getLogger(__name__)
@ -228,3 +229,60 @@ def init_email_client(options, **kwargs):
encoding=options.email_encoding,
**kwargs
)
def add_sftp_opts(parser):
""" Add SFTP options to argpase.ArgumentParser """
sftp_opts = parser.add_argument_group("SFTP options")
sftp_opts.add_argument(
'-H', '--sftp-host',
action="store",
type=str,
dest="sftp_host",
help="SFTP Host (default: localhost)",
default='localhost'
)
sftp_opts.add_argument(
'--sftp-port',
action="store",
type=int,
dest="sftp_port",
help="SFTP Port (default: 22)",
default=22
)
sftp_opts.add_argument(
'-u', '--sftp-user',
action="store",
type=str,
dest="sftp_user",
help="SFTP User"
)
sftp_opts.add_argument(
'-P', '--sftp-password',
action="store",
type=str,
dest="sftp_password",
help="SFTP Password"
)
sftp_opts.add_argument(
'--sftp-known-hosts',
action="store",
type=str,
dest="sftp_known_hosts",
help="SFTP known_hosts file path (default: ~/.ssh/known_hosts)",
default=os.path.expanduser('~/.ssh/known_hosts')
)
sftp_opts.add_argument(
'--sftp-auto-add-unknown-host-key',
action="store_true",
dest="sftp_auto_add_unknown_host_key",
help="Auto-add unknown SSH host key"
)
return sftp_opts

View file

@ -0,0 +1,69 @@
# -*- coding: utf-8 -*-
""" Test SFTP client """
import tempfile
import logging
import sys
import os
import getpass
from mylib.sftp import SFTPClient
from mylib.scripts.helpers import get_opts_parser, add_sftp_opts
from mylib.scripts.helpers import init_logging
log = logging.getLogger('mylib.scripts.sftp_test')
def main(argv=None): # pylint: disable=too-many-locals,too-many-statements
""" Script main """
if argv is None:
argv = sys.argv[1:]
# Options parser
parser = get_opts_parser(just_try=True)
add_sftp_opts(parser)
test_opts = parser.add_argument_group('Test SFTP options')
test_opts.add_argument(
'-p', '--remote-upload-path',
action="store",
type=str,
dest="upload_path",
help="Remote upload path (default: on remote initial connection directory)",
)
options = parser.parse_args()
# Initialize logs
init_logging(options, 'Test SFTP client')
if options.sftp_user and not options.sftp_password:
options.sftp_password = getpass.getpass('Please enter SFTP password: ')
log.info('Initialize Email client')
sftp = SFTPClient(options=options, just_try=options.just_try)
sftp.connect()
log.debug('Create tempory file')
tmp_file = tempfile.NamedTemporaryFile() # pylint: disable=consider-using-with
log.debug('Temporary file path: "%s"', tmp_file.name)
tmp_file.write(b'Juste un test.')
log.debug(
'Upload file %s to SFTP server (in %s)', tmp_file.name,
options.upload_path if options.upload_path else "remote initial connection directory")
if not sftp.upload_file(tmp_file.name, options.upload_path):
log.error('Fail to upload test file on SFTP server')
else:
log.info('Test file uploaded on SFTP server')
remote_filepath = (
os.path.join(options.upload_path, os.path.basename(tmp_file.name))
if options.upload_path else os.path.basename(tmp_file.name)
)
if sftp.remove_file(remote_filepath):
log.info('Test file removed on SFTP server')
else:
log.error('Fail to remove test file on SFTP server')
sftp.close()

133
mylib/sftp.py Normal file
View file

@ -0,0 +1,133 @@
# -*- coding: utf-8 -*-
""" SFTP client """
import logging
import os
from paramiko import SSHClient, AutoAddPolicy, SFTPAttributes
from mylib.config import ConfigurableObject
from mylib.config import BooleanOption
from mylib.config import IntegerOption
from mylib.config import PasswordOption
from mylib.config import StringOption
log = logging.getLogger(__name__)
class SFTPClient(ConfigurableObject):
"""
SFTP client
This class abstract all interactions with the SFTP server.
"""
_config_name = 'sftp'
_config_comment = 'SFTP'
_defaults = {
'host': 'localhost',
'port': 22,
'user': None,
'password': None,
'known_hosts': os.path.expanduser('~/.ssh/known_hosts'),
'auto_add_unknown_host_key': False,
'just_try': False,
}
ssh_client = None
sftp_client = None
initial_directory = None
def configure(self, just_try=True, ** kwargs): # pylint: disable=arguments-differ
""" Configure options on registered mylib.Config object """
section = super().configure(**kwargs)
section.add_option(
StringOption, 'host', default=self._defaults['host'],
comment='SFTP server hostname/IP address')
section.add_option(
IntegerOption, 'port', default=self._defaults['port'],
comment='SFTP server port')
section.add_option(
StringOption, 'user', default=self._defaults['user'],
comment='SFTP authentication username')
section.add_option(
PasswordOption, 'password', default=self._defaults['password'],
comment='SFTP authentication password (set to "keyring" to use XDG keyring)',
username_option='user', keyring_value='keyring')
section.add_option(
StringOption, 'known_hosts', default=self._defaults['known_hosts'],
comment='SFTP known_hosts filepath')
section.add_option(
BooleanOption, 'auto_add_unknown_host_key',
default=self._defaults['auto_add_unknown_host_key'],
comment='Auto add unknown host key')
if just_try:
section.add_option(
BooleanOption, 'just_try', default=self._defaults['just_try'],
comment='Just-try mode: do not really send emails')
return section
def initialize(self, loaded_config=None):
""" Configuration initialized hook """
super().__init__(loaded_config=loaded_config)
def connect(self):
""" Connect to SFTP server """
if self.ssh_client:
return
host = self._get_option('host')
port = self._get_option('port')
log.info("Connect to SFTP server %s:%d", host, port)
self.ssh_client = SSHClient()
if self._get_option('known_hosts'):
self.ssh_client.load_host_keys(self._get_option('known_hosts'))
if self._get_option('auto_add_unknown_host_key'):
log.debug('Set missing host key policy to auto-add')
self.ssh_client.set_missing_host_key_policy(AutoAddPolicy())
self.ssh_client.connect(
host, port=port,
username=self._get_option('user'),
password=self._get_option('password')
)
self.sftp_client = self.ssh_client.open_sftp()
self.initial_directory = self.sftp_client.getcwd()
if self.initial_directory:
log.debug("Initial remote directory: '%s'", self.initial_directory)
else:
log.debug("Fail to retreive remote directory, use empty string instead")
self.initial_directory = ""
def upload_file(self, filepath, remote_directory=None):
""" Upload a file on SFTP server """
self.connect()
remote_filepath = os.path.join(
remote_directory if remote_directory else self.initial_directory,
os.path.basename(filepath)
)
log.debug("Upload file '%s' to '%s'", filepath, remote_filepath)
if self._get_option('just_try'):
log.debug(
"Just-try mode: do not really upload file '%s' to '%s'",
filepath, remote_filepath)
return True
result = self.sftp_client.put(filepath, remote_filepath)
return isinstance(result, SFTPAttributes)
def remove_file(self, filepath):
""" Remove a file on SFTP server """
self.connect()
log.debug("Remove file '%s'", filepath)
if self._get_option('just_try'):
log.debug("Just - try mode: do not really remove file '%s'", filepath)
return True
self.sftp_client.remove(filepath)
return True
def close(self):
""" Close SSH/SFTP connection """
log.debug("Close connection")
self.ssh_client.close()

View file

@ -5,7 +5,7 @@ from setuptools import find_packages
from setuptools import setup
extras_require={
extras_require = {
'dev': [
'pytest',
'mocker',
@ -34,6 +34,9 @@ extras_require={
'mysql': [
'mysqlclient',
],
'sftp': [
'paramiko',
],
}
install_requires = ['progressbar']
@ -65,6 +68,7 @@ setup(
'mylib-test-pbar = mylib.scripts.pbar_test:main',
'mylib-test-report = mylib.scripts.report_test:main',
'mylib-test-ldap = mylib.scripts.ldap_test:main',
'mylib-test-sftp = mylib.scripts.sftp_test:main',
],
},
)