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