python-mylib/mylib/sftp.py

""" SFTP client """

import logging
import os

from paramiko import AutoAddPolicy, SFTPAttributes, SSHClient

from mylib.config import (
    BooleanOption,
    ConfigurableObject,
    IntegerOption,
    PasswordOption,
    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

    # pylint: disable=arguments-differ,arguments-renamed
    def configure(self, **kwargs):
        """Configure options on registered mylib.Config object"""
        section = super().configure(
            just_try=kwargs.pop("just_try", True),
            just_try_help=kwargs.pop(
                "just_try_help", "Just-try mode: do not really make change on remote SFTP host"
            ),
            **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",
        )

        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 retrieve 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("Retrieve 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._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._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()