# 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"