python-mylib/mylib/sftp.py

# -*- 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 get_file(self, remote_filepath, local_filepath):
        """ Retrieve a file from SFTP server """
        self.connect()
        log.debug("Retreive file '%s' to '%s'", remote_filepath, local_filepath)
        return self.sftp_client.get(remote_filepath, local_filepath) is None

    def open_file(self, remote_filepath, mode='r'):
        """ Remotly open a file on SFTP server """
        self.connect()
        log.debug("Remotly open file '%s'", remote_filepath)
        return self.sftp_client.open(remote_filepath, mode=mode)

    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
        return self.sftp_client.remove(filepath) is None

    def close(self):
        """ Close SSH/SFTP connection """
        log.debug("Close connection")
        self.ssh_client.close()