diff --git a/mylib/scripts/helpers.py b/mylib/scripts/helpers.py index e5388a0..666ce28 100644 --- a/mylib/scripts/helpers.py +++ b/mylib/scripts/helpers.py @@ -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 diff --git a/mylib/scripts/sftp_test.py b/mylib/scripts/sftp_test.py new file mode 100644 index 0000000..9cdf5dd --- /dev/null +++ b/mylib/scripts/sftp_test.py @@ -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() diff --git a/mylib/sftp.py b/mylib/sftp.py new file mode 100644 index 0000000..c5db741 --- /dev/null +++ b/mylib/sftp.py @@ -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() diff --git a/setup.py b/setup.py index f60fe39..41d2f97 100644 --- a/setup.py +++ b/setup.py @@ -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'] @@ -48,7 +51,7 @@ setup( version=version, description='A set of helpers small libs to make common tasks easier in my script development', classifiers=[ - 'Programming Language :: Python', + 'Programming Language :: Python', ], install_requires=install_requires, extras_require=extras_require, @@ -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', ], }, )