337 lines
14 KiB
Python
337 lines
14 KiB
Python
# SPDX-License-Identifier: BSD-2-Clause
|
|
# SPDX-FileCopyrightText: 2023 KDE e.V.
|
|
# SPDX-FileContributor: Ingo Klöcker <dev@ingo-kloecker.de>
|
|
|
|
import stat
|
|
import time
|
|
from unittest.mock import create_autospec, Mock
|
|
|
|
import paramiko
|
|
import pytest
|
|
|
|
import sftpnotary.sftp
|
|
from sftpnotary.sftp import SFTPClient, SFTPError
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_paramiko(monkeypatch):
|
|
monkeypatch.setattr(sftpnotary.sftp, "SSHClient", create_autospec(paramiko.SSHClient))
|
|
monkeypatch.setattr(sftpnotary.sftp.RSAKey, "from_private_key_file", Mock(return_value="-----RSA Key"))
|
|
monkeypatch.setattr(sftpnotary.sftp.RSAKey, "from_private_key", Mock(side_effect=lambda x: x.read()))
|
|
|
|
|
|
def createSFTPAttributes(**kwargs):
|
|
attrs = paramiko.SFTPAttributes()
|
|
attrs.st_mode = 0
|
|
for k, v in kwargs.items():
|
|
setattr(attrs, k, v)
|
|
return attrs
|
|
|
|
|
|
class TestSFTP:
|
|
def test_connect_and_cleanup(self):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path", knownHostsFile="/knownHostsFile", port=-1, username="username")
|
|
internalSshClient = sftp._sshClient
|
|
internalSftpClient = sftp._sftpClient
|
|
assert internalSshClient is not None
|
|
assert internalSftpClient is not None
|
|
internalSshClient.load_system_host_keys.assert_called_once_with("/knownHostsFile")
|
|
internalSshClient.set_missing_host_key_policy.assert_not_called()
|
|
internalSshClient.connect.assert_called_once_with("localhost", port=-1, username="username")
|
|
internalSftpClient.chdir.assert_called_once_with("base/path")
|
|
|
|
internalSftpClient.close.assert_called_once()
|
|
assert sftp._sftpClient is None
|
|
internalSshClient.close.assert_called_once()
|
|
assert sftp._sshClient is None
|
|
|
|
def test_connect_accepts_missing_host_keys_if_known_host_file_not_given(self):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSshClient = sftp._sshClient
|
|
internalSshClient.set_missing_host_key_policy.assert_called_once()
|
|
assert isinstance(
|
|
internalSshClient.set_missing_host_key_policy.call_args.args[0], paramiko.client.AutoAddPolicy
|
|
)
|
|
|
|
def test_connect_with_rsa_key_file(self):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path", rsaKeyFile="/rsaKey")
|
|
internalSshClient = sftp._sshClient
|
|
internalSshClient.connect.assert_called_once_with("localhost", pkey="-----RSA Key", allow_agent=False)
|
|
|
|
def test_connect_with_rsa_key_from_environment_variable(self, monkeypatch):
|
|
monkeypatch.setenv("RSA_KEY", "-----RSA Key from env. var.")
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path", rsaKeyFile="env:RSA_KEY")
|
|
internalSshClient = sftp._sshClient
|
|
internalSshClient.connect.assert_called_once_with(
|
|
"localhost", pkey="-----RSA Key from env. var.", allow_agent=False
|
|
)
|
|
|
|
def test_connect_with_rsa_key_file_from_environment_variable(self, monkeypatch):
|
|
monkeypatch.setenv("RSA_KEY", "/rsaKeyPath")
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path", rsaKeyFile="env:RSA_KEY")
|
|
internalSshClient = sftp._sshClient
|
|
internalSshClient.connect.assert_called_once_with("localhost", pkey="-----RSA Key", allow_agent=False)
|
|
|
|
def test_reconnect(self):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect(
|
|
"localhost",
|
|
"base/path",
|
|
knownHostsFile="/knownHostsFile",
|
|
rsaKeyFile="/rsaKey",
|
|
port=-1,
|
|
username="username",
|
|
)
|
|
oldSshClient = sftp._sshClient
|
|
oldSshClient.reset_mock()
|
|
oldSftpClient = sftp._sftpClient
|
|
oldSftpClient.reset_mock()
|
|
|
|
sftp.reconnect()
|
|
|
|
oldSftpClient.close.assert_called_once()
|
|
oldSshClient.close.assert_called_once()
|
|
newSshClient = sftp._sshClient
|
|
newSftpClient = sftp._sftpClient
|
|
newSshClient.connect.assert_called_once_with(
|
|
"localhost", port=-1, username="username", pkey="-----RSA Key", allow_agent=False
|
|
)
|
|
newSftpClient.chdir.assert_called_once_with("base/path")
|
|
|
|
def test_reconnect_retries_multiple_times(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
|
|
monkeypatch.setattr(sftpnotary.sftp, "SSHClient", Mock(side_effect=sftpnotary.sftp.SSHException))
|
|
monkeypatch.setattr(time, "sleep", Mock())
|
|
|
|
with pytest.raises(SFTPError):
|
|
sftp.reconnect(maxTries=10)
|
|
|
|
assert sftpnotary.sftp.SSHClient.call_count == 10
|
|
assert time.sleep.call_args_list == [
|
|
((1.0,),),
|
|
((2.0,),),
|
|
((4.0,),),
|
|
((8.0,),),
|
|
((16.0,),),
|
|
((32.0,),),
|
|
((60,),),
|
|
((60,),),
|
|
((60,),),
|
|
]
|
|
|
|
def test_download(self):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
sftp.download("remote/path", "local/path")
|
|
|
|
internalSftpClient.get.assert_called_once()
|
|
assert internalSftpClient.get.call_args.args[0] == "remote/path"
|
|
assert internalSftpClient.get.call_args.args[1] == "local/path"
|
|
|
|
def test_download_fails(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
monkeypatch.setattr(internalSftpClient, "get", Mock(side_effect=Exception))
|
|
monkeypatch.setattr(sftpnotary.sftp.Path, "unlink", Mock())
|
|
with pytest.raises(SFTPError):
|
|
sftp.download("remote/path", "local/path")
|
|
|
|
sftpnotary.sftp.Path.unlink.assert_called_once_with(missing_ok=True)
|
|
|
|
def test_downloadBytes(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
def getfo_mock(_, fl, **__):
|
|
fl.write(b"BYTES")
|
|
|
|
monkeypatch.setattr(internalSftpClient, "getfo", Mock(side_effect=getfo_mock))
|
|
result = sftp.downloadBytes("remote/path")
|
|
|
|
assert result == b"BYTES"
|
|
internalSftpClient.getfo.assert_called_once()
|
|
assert internalSftpClient.getfo.call_args.args[0] == "remote/path"
|
|
|
|
def test_downloadBytes_fails(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
monkeypatch.setattr(internalSftpClient, "getfo", Mock(side_effect=Exception))
|
|
with pytest.raises(SFTPError):
|
|
sftp.downloadBytes("remote/path")
|
|
|
|
def test_upload(self):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
sftp.upload("local/path", "remote/path")
|
|
|
|
internalSftpClient.put.assert_called_once()
|
|
assert internalSftpClient.put.call_args.args[0] == "local/path"
|
|
assert internalSftpClient.put.call_args.args[1] == "remote/path"
|
|
|
|
def test_upload_retried_up_to_3_times_on_EOFError(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
monkeypatch.setattr(sftp, "reconnect", Mock())
|
|
monkeypatch.setattr(internalSftpClient, "put", Mock(side_effect=(EOFError, None)))
|
|
sftp.upload("local/path", "remote/path")
|
|
assert sftp.reconnect.called
|
|
assert internalSftpClient.put.call_count == 2
|
|
|
|
sftp.reconnect.reset_mock()
|
|
internalSftpClient.put.reset_mock()
|
|
internalSftpClient.put.side_effect = (EOFError, EOFError, None)
|
|
sftp.upload("local/path", "remote/path")
|
|
assert internalSftpClient.put.call_count == 3
|
|
|
|
sftp.reconnect.reset_mock()
|
|
internalSftpClient.put.reset_mock()
|
|
internalSftpClient.put.side_effect = EOFError
|
|
with pytest.raises(SFTPError):
|
|
sftp.upload("local/path", "remote/path")
|
|
assert internalSftpClient.put.call_count == 3
|
|
|
|
def test_upload_not_retried_on_other_errors(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
monkeypatch.setattr(sftp, "reconnect", Mock())
|
|
monkeypatch.setattr(internalSftpClient, "put", Mock(side_effect=Exception))
|
|
with pytest.raises(SFTPError):
|
|
sftp.upload("local/path", "remote/path")
|
|
assert not sftp.reconnect.called
|
|
assert internalSftpClient.put.call_count == 1
|
|
|
|
def test_uploadBytes(self):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
sftp.uploadBytes(b"BYTES", "remote/path")
|
|
|
|
internalSftpClient.putfo.assert_called_once()
|
|
assert internalSftpClient.putfo.call_args.args[1] == "remote/path"
|
|
|
|
def test_uploadBytes_retried_up_to_3_times_on_EOFError(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
monkeypatch.setattr(sftp, "reconnect", Mock())
|
|
monkeypatch.setattr(internalSftpClient, "putfo", Mock(side_effect=(EOFError, None)))
|
|
sftp.uploadBytes(b"BYTES", "remote/path")
|
|
assert sftp.reconnect.called
|
|
assert internalSftpClient.putfo.call_count == 2
|
|
|
|
sftp.reconnect.reset_mock()
|
|
internalSftpClient.putfo.reset_mock()
|
|
internalSftpClient.putfo.side_effect = (EOFError, EOFError, None)
|
|
sftp.uploadBytes(b"BYTES", "remote/path")
|
|
assert internalSftpClient.putfo.call_count == 3
|
|
|
|
sftp.reconnect.reset_mock()
|
|
internalSftpClient.putfo.reset_mock()
|
|
internalSftpClient.putfo.side_effect = EOFError
|
|
with pytest.raises(SFTPError):
|
|
sftp.uploadBytes(b"BYTES", "remote/path")
|
|
assert internalSftpClient.putfo.call_count == 3
|
|
|
|
def test_uploadBytes_not_retried_on_other_errors(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
monkeypatch.setattr(sftp, "reconnect", Mock())
|
|
monkeypatch.setattr(internalSftpClient, "putfo", Mock(side_effect=Exception))
|
|
with pytest.raises(SFTPError):
|
|
sftp.uploadBytes(b"BYTES", "remote/path")
|
|
assert not sftp.reconnect.called
|
|
assert internalSftpClient.putfo.call_count == 1
|
|
|
|
def test_recursiveRemoveDirectory(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
def listdir_attr_mock(path):
|
|
if str(path) == "folder":
|
|
return [
|
|
createSFTPAttributes(filename="file1"),
|
|
createSFTPAttributes(filename="file2"),
|
|
createSFTPAttributes(filename="subfolder1", st_mode=stat.S_IFDIR),
|
|
]
|
|
elif str(path) == "folder/subfolder1":
|
|
return [
|
|
createSFTPAttributes(filename="file1.1"),
|
|
createSFTPAttributes(filename="file1.2"),
|
|
]
|
|
raise Exception(f"Called with unexpected path {path}")
|
|
|
|
monkeypatch.setattr(internalSftpClient, "listdir_attr", Mock(side_effect=listdir_attr_mock))
|
|
monkeypatch.setattr(internalSftpClient, "remove", Mock())
|
|
monkeypatch.setattr(internalSftpClient, "rmdir", Mock())
|
|
|
|
sftp.recursiveRemoveDirectory("folder")
|
|
|
|
assert internalSftpClient.remove.call_count == 4
|
|
assert internalSftpClient.remove.call_args_list == [
|
|
(("folder/file1",),),
|
|
(("folder/file2",),),
|
|
(("folder/subfolder1/file1.1",),),
|
|
(("folder/subfolder1/file1.2",),),
|
|
]
|
|
assert internalSftpClient.rmdir.call_count == 2
|
|
assert internalSftpClient.rmdir.call_args_list == [
|
|
(("folder/subfolder1",),),
|
|
(("folder",),),
|
|
]
|
|
|
|
def test_forwarded_functions(self, monkeypatch):
|
|
with SFTPClient() as sftp:
|
|
sftp.connect("localhost", "base/path")
|
|
internalSftpClient = sftp._sftpClient
|
|
|
|
monkeypatch.setattr(internalSftpClient, "listdir", Mock(return_value="listdir"))
|
|
assert sftp.listdir("remote/path") == "listdir"
|
|
|
|
monkeypatch.setattr(internalSftpClient, "listdir_attr", Mock(return_value="listdir_attr"))
|
|
assert sftp.listdir_attr("remote/path") == "listdir_attr"
|
|
|
|
monkeypatch.setattr(internalSftpClient, "listdir_iter", Mock(return_value="listdir_iter"))
|
|
assert sftp.listdir_iter("remote/path") == "listdir_iter"
|
|
|
|
monkeypatch.setattr(internalSftpClient, "lstat", Mock(return_value="lstat"))
|
|
assert sftp.lstat("remote/path") == "lstat"
|
|
|
|
monkeypatch.setattr(internalSftpClient, "mkdir", Mock(return_value="mkdir"))
|
|
assert sftp.mkdir("remote/path") == "mkdir"
|
|
|
|
monkeypatch.setattr(internalSftpClient, "posix_rename", Mock(return_value="posix_rename"))
|
|
assert sftp.posix_rename("remote/oldpath", "remote/newpath") == "posix_rename"
|
|
|
|
monkeypatch.setattr(internalSftpClient, "remove", Mock(return_value="remove"))
|
|
assert sftp.remove("remote/path") == "remove"
|
|
|
|
monkeypatch.setattr(internalSftpClient, "rmdir", Mock(return_value="rmdir"))
|
|
assert sftp.rmdir("remote/path") == "rmdir"
|
|
|
|
monkeypatch.setattr(internalSftpClient, "stat", Mock(return_value="stat"))
|
|
assert sftp.stat("remote/path") == "stat"
|