mirror of http://git.sairate.top/sairate/doc.git
398 lines
14 KiB
Python
398 lines
14 KiB
Python
import errno
|
|
import logging
|
|
import stat
|
|
from pathlib import PurePosixPath
|
|
from types import TracebackType
|
|
from typing import Any, cast, Generator, Iterable, Optional
|
|
|
|
import paramiko
|
|
from cerulean.file_system import FileSystem
|
|
from cerulean.file_system_impl import FileSystemImpl
|
|
from cerulean.path import AbstractPath, EntryType, Path, Permission
|
|
from cerulean.ssh_terminal import SshTerminal
|
|
from cerulean.util import BaseExceptionType
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SftpFileSystem(FileSystemImpl):
|
|
"""A FileSystem implementation that connects to an SFTP server.
|
|
|
|
SftpFileSystem supports the / operation:
|
|
|
|
.. code-block:: python
|
|
|
|
fs / 'path'
|
|
|
|
which produces a :class:`Path`, through which you can do things \
|
|
with the remote files.
|
|
|
|
It is also a context manager, so that you can (and should!) use it \
|
|
with a ``with`` statement, which will ensure that the connection \
|
|
is closed when you are done with the it. Alternatively, you can \
|
|
call :meth:`close` to close the connection.
|
|
|
|
If `own_term` is True, this class assumes that it owns the terminal \
|
|
you gave it, and that it is responsible for closing it when it's \
|
|
done with it. If you share an SshTerminal between an SftpFileSystem \
|
|
and a scheduler, or use the terminal directly yourself, then you \
|
|
want to use False here, and close the terminal yourself when you \
|
|
don't need it any more.
|
|
|
|
Args:
|
|
terminal: The terminal to connect through.
|
|
own_term: Whether to close the terminal when the file system \
|
|
is closed.
|
|
|
|
"""
|
|
def __init__(self, terminal: SshTerminal, own_term: bool = False) -> None:
|
|
self.__terminal = terminal
|
|
self.__own_term = own_term
|
|
self.__ensure_sftp(True)
|
|
self.__sftp2 = None # type: Optional[paramiko.SFTPClient]
|
|
self.__max_tries = 3
|
|
|
|
def __enter__(self) -> 'SftpFileSystem':
|
|
"""Enter context manager."""
|
|
return self
|
|
|
|
def __exit__(self, exc_type: Optional[BaseExceptionType],
|
|
exc_value: Optional[BaseException],
|
|
traceback: Optional[TracebackType]) -> None:
|
|
"""Exit context manager."""
|
|
if self.__own_term:
|
|
self.close()
|
|
|
|
def close(self) -> None:
|
|
self.__sftp.close()
|
|
logger.info('Disconnected from SFTP server')
|
|
|
|
def __eq__(self, other: Any) -> bool:
|
|
if not isinstance(other, FileSystem):
|
|
return NotImplemented
|
|
if isinstance(other, SftpFileSystem):
|
|
return self.__terminal == other.__terminal
|
|
return False
|
|
|
|
def root(self) -> Path:
|
|
return Path(self, PurePosixPath('/'))
|
|
|
|
def __truediv__(self, segment: str) -> Path:
|
|
return Path(self, PurePosixPath('/' + segment.strip('/')))
|
|
|
|
def _supports(self, feature: str) -> bool:
|
|
if feature not in self._features:
|
|
raise ValueError('Invalid argument for "feature"')
|
|
return True
|
|
|
|
def _exists(self, path: AbstractPath) -> bool:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
try:
|
|
self.__sftp.stat(str(lpath))
|
|
return True
|
|
except IOError:
|
|
return False
|
|
|
|
def _mkdir(
|
|
self, path: AbstractPath, mode: Optional[int] = None,
|
|
parents: bool = False, exists_ok: bool = False) -> None:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
if parents:
|
|
for parent in reversed(lpath.parents):
|
|
if not self._exists(parent):
|
|
self.__sftp.mkdir(str(parent))
|
|
if self._exists(lpath):
|
|
if not exists_ok:
|
|
raise FileExistsError(
|
|
'File {} exists and exists_ok was False'.format(lpath))
|
|
else:
|
|
return
|
|
|
|
if mode is not None:
|
|
self.__sftp.mkdir(str(lpath), mode)
|
|
else:
|
|
self.__sftp.mkdir(str(lpath))
|
|
|
|
def _iterdir(
|
|
self, path: AbstractPath) -> Generator[PurePosixPath, None, None]:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
# Note: we're not using listdir_iter here, because it hangs:
|
|
# https://github.com/paramiko/paramiko/issues/1171
|
|
try:
|
|
for entry in self.__sftp.listdir(str(lpath)):
|
|
yield lpath / entry
|
|
except OSError as e:
|
|
# Paramiko omits the filename, which breaks Path.walk()
|
|
# so add it back in here.
|
|
raise OSError(e.errno, e.strerror, str(lpath))
|
|
|
|
def _rmdir(self, path: AbstractPath, recursive: bool = False) -> None:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
if not self._exists(lpath):
|
|
return
|
|
|
|
if not self._is_dir(lpath):
|
|
raise RuntimeError("Path must refer to a directory")
|
|
|
|
if recursive:
|
|
for entry in self.__sftp.listdir_attr(str(lpath)):
|
|
entry_path = lpath / entry.filename
|
|
if self._is_symlink(entry_path):
|
|
self.__sftp.unlink(str(entry_path))
|
|
elif self._is_dir(entry_path):
|
|
self._rmdir(entry_path, True)
|
|
else:
|
|
self.__sftp.unlink(str(entry_path))
|
|
|
|
self.__sftp.rmdir(str(lpath))
|
|
|
|
def _touch(self, path: AbstractPath) -> None:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
with self.__sftp.file(str(lpath), 'a'):
|
|
pass
|
|
|
|
def _streaming_read(
|
|
self, path: AbstractPath) -> Generator[bytes, None, None]:
|
|
# Buffer size vs. speed (MB/s) against localhost
|
|
# up down local
|
|
# 8k 33 56 159
|
|
# 16k 52 56 145
|
|
# 24k 66 57 150
|
|
# 32k 24 57 149
|
|
# 2M 24 48
|
|
# scp 120 110
|
|
# cp 172
|
|
def ensure_sftp2(self: 'SftpFileSystem') -> None:
|
|
if self.__sftp2 is None:
|
|
self.__sftp2 = self.__terminal._get_downstream_sftp_client()
|
|
else:
|
|
try:
|
|
self.__sftp2.lstat('/')
|
|
except OSError as e:
|
|
if 'Socket is closed' in str(e):
|
|
self.__sftp2 = self.__terminal._get_downstream_sftp_client()
|
|
else:
|
|
raise
|
|
|
|
channel = self.__sftp2.get_channel()
|
|
if not (channel and channel.get_transport().is_active()):
|
|
self.__sftp2 = self.__terminal._get_downstream_sftp_client()
|
|
|
|
lpath = cast(PurePosixPath, path)
|
|
ensure_sftp2(self)
|
|
try:
|
|
size = self._size(path)
|
|
with self.__sftp2.file(str(lpath), 'rb') as f: # type: ignore
|
|
f.prefetch(size)
|
|
data = f.read(24576)
|
|
while len(data) > 0:
|
|
yield data
|
|
data = f.read(24576)
|
|
except paramiko.SSHException as e:
|
|
if 'Server connection dropped' in str(e):
|
|
raise ConnectionError(e)
|
|
else:
|
|
raise e
|
|
|
|
def _streaming_write(
|
|
self, path: AbstractPath, data: Iterable[bytes]) -> None:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
try:
|
|
with self.__sftp.file(str(lpath), 'wb') as f:
|
|
f.set_pipelined(True)
|
|
for chunk in data:
|
|
f.write(chunk)
|
|
except paramiko.SSHException as e:
|
|
if 'Server connection dropped' in str(e):
|
|
raise ConnectionError(e)
|
|
else:
|
|
raise e
|
|
|
|
def _rename(self, path: AbstractPath, target: AbstractPath) -> None:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
ltarget = cast(PurePosixPath, target)
|
|
self.__sftp.posix_rename(str(lpath), str(ltarget))
|
|
|
|
def _unlink(self, path: AbstractPath) -> None:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
self.__sftp.unlink(str(lpath))
|
|
|
|
def _is_dir(self, path: AbstractPath) -> bool:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
try:
|
|
mode = self.__stat(lpath).st_mode
|
|
return mode is not None and bool(stat.S_ISDIR(mode))
|
|
except FileNotFoundError:
|
|
return False
|
|
|
|
def _is_file(self, path: AbstractPath) -> bool:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
try:
|
|
mode = self.__stat(lpath).st_mode
|
|
return mode is not None and bool(stat.S_ISREG(mode))
|
|
except FileNotFoundError:
|
|
return False
|
|
|
|
def _is_symlink(self, path: AbstractPath) -> bool:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
try:
|
|
mode = self.__lstat(lpath).st_mode
|
|
return mode is not None and bool(stat.S_ISLNK(mode))
|
|
except FileNotFoundError:
|
|
return False
|
|
|
|
def _entry_type(self, path: AbstractPath) -> EntryType:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
mode_to_type = [(stat.S_ISDIR, EntryType.DIRECTORY),
|
|
(stat.S_ISREG, EntryType.FILE),
|
|
(stat.S_ISLNK, EntryType.SYMBOLIC_LINK),
|
|
(stat.S_ISCHR, EntryType.CHARACTER_DEVICE),
|
|
(stat.S_ISBLK, EntryType.BLOCK_DEVICE),
|
|
(stat.S_ISFIFO, EntryType.FIFO),
|
|
(stat.S_ISSOCK, EntryType.SOCKET)]
|
|
|
|
try:
|
|
mode = self.__lstat(lpath).st_mode
|
|
except IOError:
|
|
raise OSError(errno.ENOENT, 'No such file or directory',
|
|
str(lpath))
|
|
|
|
for predicate, result in mode_to_type:
|
|
if mode is not None and predicate(mode):
|
|
return result
|
|
raise RuntimeError('Object is of unknown type, please report a'
|
|
'Cerulean bug')
|
|
|
|
def _size(self, path: AbstractPath) -> int:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
size = self.__stat(lpath).st_size
|
|
if size is None:
|
|
raise RuntimeError('Server did not return size')
|
|
return size
|
|
|
|
def _uid(self, path: AbstractPath) -> int:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
uid = self.__stat(lpath).st_uid
|
|
if uid is None:
|
|
raise RuntimeError('Server did not return a UID')
|
|
return uid
|
|
|
|
def _gid(self, path: AbstractPath) -> int:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
gid = self.__stat(lpath).st_gid
|
|
if gid is None:
|
|
raise RuntimeError('Server did not return a GID')
|
|
return gid
|
|
|
|
def _has_permission(
|
|
self, path: AbstractPath, permission: Permission) -> bool:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
mode = self.__stat(lpath).st_mode
|
|
if mode is None:
|
|
raise RuntimeError('Server did not return file mode')
|
|
return bool(mode & permission.value)
|
|
|
|
def _set_permission(
|
|
self, path: AbstractPath, permission: Permission,
|
|
value: bool = True) -> None:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
mode = self.__stat(lpath).st_mode
|
|
if mode is None:
|
|
raise RuntimeError('Server did not return file mode')
|
|
if value:
|
|
mode = mode | permission.value
|
|
else:
|
|
mode = mode & ~permission.value
|
|
self._chmod(lpath, mode)
|
|
|
|
def _chmod(self, path: AbstractPath, mode: int) -> None:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
self.__sftp.chmod(str(lpath), mode)
|
|
|
|
def _symlink_to(self, path: AbstractPath, target: AbstractPath) -> None:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
ltarget = cast(PurePosixPath, target)
|
|
self.__sftp.symlink(str(ltarget), str(lpath))
|
|
|
|
def _readlink(self, path: AbstractPath, recursive: bool) -> Path:
|
|
self.__ensure_sftp()
|
|
lpath = cast(PurePosixPath, path)
|
|
if recursive:
|
|
# SFTP's normalize() raises if there's a link loop or a \
|
|
# non-existing target, which we don't want, so we use \
|
|
# our own algorithm.
|
|
max_iter = 32
|
|
cur_path = lpath
|
|
iter_count = 0
|
|
while self._is_symlink(cur_path) and iter_count < max_iter:
|
|
target_str = self.__sftp.readlink(str(cur_path))
|
|
if target_str is None:
|
|
raise RuntimeError('Server error while reading link')
|
|
target = PurePosixPath(target_str)
|
|
if not target.is_absolute():
|
|
target = cur_path.parent / target
|
|
cur_path = target
|
|
iter_count += 1
|
|
|
|
if iter_count == max_iter:
|
|
raise RuntimeError('Too many symbolic links detected')
|
|
|
|
target = PurePosixPath(self.__sftp.normalize(str(path)))
|
|
else:
|
|
target_str = self.__sftp.readlink(str(path))
|
|
if target_str is None:
|
|
raise RuntimeError('Server error while reading link')
|
|
target = PurePosixPath(target_str)
|
|
if not target.is_absolute():
|
|
target = lpath.parent / target
|
|
|
|
return Path(self, target)
|
|
|
|
def __lstat(self, path: PurePosixPath) -> paramiko.SFTPAttributes:
|
|
return self.__sftp.lstat(str(path))
|
|
|
|
def __stat(self, path: PurePosixPath) -> paramiko.SFTPAttributes:
|
|
return self.__sftp.stat(str(path))
|
|
|
|
def __ensure_sftp(self, first: bool = False) -> None:
|
|
if first:
|
|
logger.info('Connecting to SFTP server')
|
|
self.__sftp = self.__terminal._get_sftp_client()
|
|
logger.info('Connected to SFTP server')
|
|
else:
|
|
try:
|
|
self.__sftp.lstat('/')
|
|
except OSError as e:
|
|
if 'Socket is closed' in str(e):
|
|
logger.info('Reconnecting to SFTP server')
|
|
self.__sftp = self.__terminal._get_sftp_client()
|
|
logger.info('Connected to SFTP server')
|
|
else:
|
|
raise
|
|
|
|
channel = self.__sftp.get_channel()
|
|
if not (channel and channel.get_transport().is_active()):
|
|
logger.info('Reconnecting to SFTP server')
|
|
self.__sftp = self.__terminal._get_sftp_client()
|
|
logger.info('Connected to SFTP server')
|