320 lines
12 KiB
Python
320 lines
12 KiB
Python
# SPDX-License-Identifier: BSD-2-Clause
|
|
# SPDX-FileCopyrightText: 2023 KDE e.V.
|
|
# SPDX-FileContributor: Ingo Klöcker <dev@ingo-kloecker.de>
|
|
|
|
import io
|
|
import logging
|
|
import os
|
|
import stat
|
|
import time
|
|
from functools import wraps
|
|
from pathlib import Path, PurePath
|
|
from typing import Optional, Union
|
|
|
|
from paramiko import RSAKey, SSHClient
|
|
from paramiko.client import AutoAddPolicy
|
|
from paramiko.common import o777
|
|
from paramiko.ssh_exception import SSHException
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class SFTPError(Exception):
|
|
pass
|
|
|
|
|
|
class ConnectionLostError(SFTPError):
|
|
pass
|
|
|
|
|
|
class ClosingContextManager:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, type, value, traceback):
|
|
self.close()
|
|
|
|
|
|
def ProgressLogger(message: str, interval: int = 1):
|
|
lastLog = time.time()
|
|
|
|
def printProgress(current: int, total: int):
|
|
nonlocal lastLog
|
|
now = time.time()
|
|
if current == total or now > lastLog + interval:
|
|
lastLog = now
|
|
log.debug(message, current, total)
|
|
|
|
return printProgress
|
|
|
|
|
|
def toSftpPath(path: Union[str, PurePath]):
|
|
if isinstance(path, PurePath):
|
|
return str(path.as_posix())
|
|
return path
|
|
|
|
|
|
def detectLostConnection(f):
|
|
@wraps(f)
|
|
def wrapper(self, *args, **kwargs):
|
|
try:
|
|
return f(self, *args, **kwargs)
|
|
except OSError:
|
|
# probably a socket error; check if the connection was closed
|
|
if self.closed:
|
|
raise ConnectionLostError("Lost connection to SFTP server.")
|
|
else:
|
|
raise
|
|
|
|
return wrapper
|
|
|
|
|
|
class SFTPClient(ClosingContextManager):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
self._connectInfo = None
|
|
self._sshClient = None
|
|
self._sftpClient = None
|
|
self.basePath = None
|
|
|
|
@property
|
|
def sftpClient(self):
|
|
if self._sftpClient is None:
|
|
self._sftpClient = self._sshClient.open_sftp()
|
|
if self.basePath is not None:
|
|
try:
|
|
path = toSftpPath(self.basePath)
|
|
self._sftpClient.chdir(path)
|
|
path = None
|
|
except FileNotFoundError:
|
|
log.debug(f"Folder '{path}' does not exist. Trying to create it ...")
|
|
if path is not None:
|
|
try:
|
|
self._sftpClient.mkdir(path)
|
|
self._sftpClient.chdir(path)
|
|
except OSError as e:
|
|
raise SFTPError(
|
|
f"Failed to create folder '{path}' on the SFTP server ({e!r}). Create it manually."
|
|
)
|
|
return self._sftpClient
|
|
|
|
def close(self):
|
|
if self._sftpClient is not None:
|
|
self._sftpClient.close()
|
|
self._sftpClient = None
|
|
if self._sshClient is not None:
|
|
self._sshClient.close()
|
|
self._sshClient = None
|
|
|
|
@property
|
|
def closed(self):
|
|
return self._sftpClient is None or self._sftpClient.get_channel()._closed
|
|
|
|
def connect(
|
|
self,
|
|
hostname: str,
|
|
basePath: Union[str, PurePath],
|
|
knownHostsFile: Optional[Union[str, Path]] = None,
|
|
rsaKeyFile: Optional[Union[str, Path]] = None,
|
|
**kwargs,
|
|
):
|
|
self._connectInfo = {
|
|
"args": (hostname, basePath),
|
|
"kwargs": {**kwargs, "knownHostsFile": knownHostsFile, "rsaKeyFile": rsaKeyFile},
|
|
}
|
|
|
|
self.basePath = basePath
|
|
|
|
connectArgs = dict(**kwargs)
|
|
if rsaKeyFile is not None:
|
|
if isinstance(rsaKeyFile, Path):
|
|
rsaKeyFile = str(rsaKeyFile)
|
|
# explicitly load the RSA key file because using SSHClient.connect with key_file fails
|
|
# (https://github.com/paramiko/paramiko/issues/1839)
|
|
if isinstance(rsaKeyFile, str) and rsaKeyFile.startswith("env:"):
|
|
envVarValue = os.environ[rsaKeyFile[4:]]
|
|
# guess whether it's a file path or an actual RSA key
|
|
if envVarValue.startswith("-----"):
|
|
# looks like a private (RSA) key
|
|
privateKey = io.StringIO(envVarValue)
|
|
connectArgs["pkey"] = RSAKey.from_private_key(privateKey)
|
|
else:
|
|
connectArgs["pkey"] = RSAKey.from_private_key_file(envVarValue)
|
|
else:
|
|
connectArgs["pkey"] = RSAKey.from_private_key_file(rsaKeyFile)
|
|
connectArgs["allow_agent"] = False # use agent only if we do not have a key
|
|
self._sshClient = SSHClient()
|
|
if isinstance(knownHostsFile, Path):
|
|
knownHostsFile = str(knownHostsFile)
|
|
self._sshClient.load_system_host_keys(knownHostsFile)
|
|
if knownHostsFile is None:
|
|
# Set "nosec B507" as this is intentional behaviour
|
|
self._sshClient.set_missing_host_key_policy(AutoAddPolicy()) # nosec B507
|
|
self._sshClient.connect(hostname, **connectArgs)
|
|
self.sftpClient
|
|
|
|
def reconnect(self, maxTries: int = -1, timeout: int = 30 * 60):
|
|
if self._connectInfo is None:
|
|
# connect() was never called
|
|
log.warning("Connect info is None. Ignoring reconnect request.")
|
|
return
|
|
|
|
# ensure that the previous connection is closed properly
|
|
self.close()
|
|
|
|
tries = 0
|
|
endTime = time.time() + timeout
|
|
maxCoolingOffPeriod = 60
|
|
coolingOffPeriod = 1.0
|
|
while True:
|
|
tries += 1
|
|
e: Optional[Exception] = None
|
|
try:
|
|
self.connect(*self._connectInfo["args"], **self._connectInfo["kwargs"])
|
|
return
|
|
except SSHException as ex:
|
|
e = ex
|
|
except OSError as ex:
|
|
e = ex
|
|
if maxTries > 0 and tries >= maxTries:
|
|
raise SFTPError(f"Failed to reconnect to the SFTP server ({e!r}). Giving up after {maxTries} retries.")
|
|
now = time.time()
|
|
if now > endTime:
|
|
raise SFTPError(f"Failed to reconnect to the SFTP server ({e!r}). Giving up after {timeout} seconds.")
|
|
coolingOffPeriod = min(coolingOffPeriod, endTime - now)
|
|
log.debug(f"Failed to reconnect to the SFTP server ({e!r}). Retrying in {coolingOffPeriod} s ...")
|
|
time.sleep(coolingOffPeriod)
|
|
coolingOffPeriod = min(coolingOffPeriod * 2, maxCoolingOffPeriod)
|
|
|
|
def download(self, remotePath: Union[str, PurePath], localPath: Union[str, Path]):
|
|
if isinstance(localPath, str):
|
|
localPath = Path(localPath)
|
|
log.info(f"Downloading file {remotePath!s} to {localPath!s} ...")
|
|
try:
|
|
self.sftpClient.get(
|
|
toSftpPath(remotePath), toSftpPath(localPath), ProgressLogger("Downloaded %s bytes of %s bytes")
|
|
)
|
|
except Exception as e:
|
|
# remove a potentially existing partially downloaded file
|
|
localPath.unlink(missing_ok=True)
|
|
raise SFTPError(f"Failed to download file '{remotePath!s}' from the SFTP server ({e!r})")
|
|
|
|
def downloadBytes(self, remotePath: Union[str, PurePath]):
|
|
log.info(f"Downloading file {remotePath!s} ...")
|
|
try:
|
|
with io.BytesIO() as dataStream:
|
|
self.sftpClient.getfo(
|
|
toSftpPath(remotePath), dataStream, callback=ProgressLogger("Downloaded %s bytes of %s bytes")
|
|
)
|
|
return dataStream.getvalue()
|
|
except Exception as e:
|
|
raise SFTPError(f"Failed to download file '{remotePath!s}' from the SFTP server ({e!r})")
|
|
|
|
def upload(self, localPath: Union[str, PurePath], remotePath: Union[str, PurePath]):
|
|
log.info(f"Uploading file {localPath!s} to {remotePath!s} ...")
|
|
|
|
maxTries = 3
|
|
tries = 0
|
|
while True:
|
|
exception = None
|
|
try:
|
|
self.sftpClient.put(
|
|
toSftpPath(localPath), toSftpPath(remotePath), ProgressLogger("Uploaded %s bytes of %s bytes")
|
|
)
|
|
return
|
|
except EOFError as e:
|
|
tries += 1
|
|
exception = e
|
|
except Exception as e:
|
|
raise SFTPError(f"Failed to upload file '{localPath!s}' to the SFTP server ({e!r})")
|
|
|
|
if tries < maxTries:
|
|
log.error("Error: Uploading file '%s' failed with %r. Retrying ...", localPath, exception)
|
|
self.reconnect()
|
|
try:
|
|
self.sftpClient.remove(toSftpPath(remotePath))
|
|
except FileNotFoundError:
|
|
pass
|
|
else:
|
|
raise SFTPError(f"Failed to upload file '{localPath!s}' to the SFTP server ({exception!r})")
|
|
|
|
def uploadBytes(self, data: bytes, remotePath: Union[str, PurePath]):
|
|
dataSize = len(data)
|
|
log.info(f"Uploading {dataSize} bytes to {remotePath!s} ...")
|
|
|
|
maxTries = 3
|
|
tries = 0
|
|
while True:
|
|
dataStream = io.BytesIO(data)
|
|
exception = None
|
|
try:
|
|
self.sftpClient.putfo(
|
|
dataStream,
|
|
toSftpPath(remotePath),
|
|
file_size=dataSize,
|
|
callback=ProgressLogger("Uploaded %s bytes of %s bytes"),
|
|
)
|
|
return
|
|
except EOFError as e:
|
|
tries += 1
|
|
exception = e
|
|
except Exception as e:
|
|
raise SFTPError(f"Failed to upload data to the SFTP server ({e!r})")
|
|
|
|
if tries < maxTries:
|
|
log.error(f"Error: Uploading data failed with {exception!r}. Retrying ...")
|
|
self.reconnect()
|
|
try:
|
|
self.sftpClient.remove(toSftpPath(remotePath))
|
|
except FileNotFoundError:
|
|
pass
|
|
else:
|
|
raise SFTPError(f"Failed to upload data to the SFTP server ({exception!r})")
|
|
|
|
@detectLostConnection
|
|
def listdir(self, path: Union[str, PurePath] = "."):
|
|
return self.sftpClient.listdir(toSftpPath(path))
|
|
|
|
@detectLostConnection
|
|
def listdir_attr(self, path: Union[str, PurePath] = "."):
|
|
return self.sftpClient.listdir_attr(toSftpPath(path))
|
|
|
|
@detectLostConnection
|
|
def listdir_iter(self, path: Union[str, PurePath] = ".", read_aheads=50):
|
|
return self.sftpClient.listdir_iter(toSftpPath(path), read_aheads=read_aheads)
|
|
|
|
@detectLostConnection
|
|
def lstat(self, path: Union[str, PurePath]):
|
|
return self.sftpClient.lstat(toSftpPath(path))
|
|
|
|
@detectLostConnection
|
|
def mkdir(self, path: Union[str, PurePath], mode=o777):
|
|
return self.sftpClient.mkdir(toSftpPath(path), mode=mode)
|
|
|
|
@detectLostConnection
|
|
def posix_rename(self, oldpath: Union[str, PurePath], newpath: Union[str, PurePath]):
|
|
return self.sftpClient.posix_rename(toSftpPath(oldpath), toSftpPath(newpath))
|
|
|
|
@detectLostConnection
|
|
def remove(self, path: Union[str, PurePath]):
|
|
return self.sftpClient.remove(toSftpPath(path))
|
|
|
|
@detectLostConnection
|
|
def rmdir(self, path: Union[str, PurePath]):
|
|
return self.sftpClient.rmdir(toSftpPath(path))
|
|
|
|
def recursiveRemoveDirectory(self, path: Union[str, PurePath]):
|
|
path = PurePath(path)
|
|
entriesAttributes = self.listdir_attr(path)
|
|
for entryAttributes in entriesAttributes:
|
|
entryPath = path / entryAttributes.filename
|
|
if stat.S_ISDIR(entryAttributes.st_mode):
|
|
self.recursiveRemoveDirectory(entryPath)
|
|
else:
|
|
self.remove(entryPath)
|
|
self.rmdir(path)
|
|
|
|
@detectLostConnection
|
|
def stat(self, path: Union[str, PurePath]):
|
|
return self.sftpClient.stat(toSftpPath(path))
|