python-mylib/mylib/sftp.py

162 lines
5.3 KiB
Python

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