Add SFTP client

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

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()

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()

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',
         ],
     },
 )