import logging
import os

import six
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from dagster_ssh.resources import SSHResource, key_from_str
from dagster_ssh.resources import ssh_resource as sshresource

from dagster import Field, ModeDefinition, execute_solid, solid
from dagster.seven import get_system_temp_directory, mock


def generate_ssh_key():
    # generate private/public key pair
    key = rsa.generate_private_key(backend=default_backend(), public_exponent=65537, key_size=2048)

    # get private key in PEM container format
    return six.ensure_str(
        key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )
    )


@mock.patch("paramiko.SSHClient")
def test_ssh_connection_with_password(ssh_mock):
    ssh_resource = SSHResource(
        remote_host="remote_host",
        remote_port=12345,
        username="username",
        password="password",
        key_file="fake.file",
        timeout=10,
        keepalive_interval=30,
        compress=True,
        no_host_key_check=False,
        allow_host_key_change=False,
        logger=logging.root.getChild("test_resources"),
    )

    with ssh_resource.get_connection():
        ssh_mock.return_value.connect.assert_called_once_with(
            hostname="remote_host",
            username="username",
            password="password",
            pkey=None,
            key_filename="fake.file",
            timeout=10,
            compress=True,
            port=12345,
            sock=None,
            look_for_keys=False,
        )


@mock.patch("paramiko.SSHClient")
def test_ssh_connection_without_password(ssh_mock):
    ssh_resource = SSHResource(
        remote_host="remote_host",
        remote_port=12345,
        username="username",
        password=None,
        timeout=10,
        key_file="fake.file",
        keepalive_interval=30,
        compress=True,
        no_host_key_check=False,
        allow_host_key_change=False,
        logger=logging.root.getChild("test_resources"),
    )

    with ssh_resource.get_connection():
        ssh_mock.return_value.connect.assert_called_once_with(
            hostname="remote_host",
            username="username",
            pkey=None,
            key_filename="fake.file",
            timeout=10,
            compress=True,
            port=12345,
            sock=None,
        )


@mock.patch("paramiko.SSHClient")
def test_ssh_connection_with_key_string(ssh_mock):
    ssh_key = generate_ssh_key()

    ssh_resource = SSHResource(
        remote_host="remote_host",
        remote_port=12345,
        username="username",
        password=None,
        timeout=10,
        key_string=six.ensure_str(ssh_key),
        keepalive_interval=30,
        compress=True,
        no_host_key_check=False,
        allow_host_key_change=False,
        logger=logging.root.getChild("test_resources"),
    )

    with ssh_resource.get_connection():
        ssh_mock.return_value.connect.assert_called_once_with(
            hostname="remote_host",
            username="username",
            key_filename=None,
            pkey=key_from_str(ssh_key),
            timeout=10,
            compress=True,
            port=12345,
            sock=None,
        )


@mock.patch("dagster_ssh.resources.SSHTunnelForwarder")
def test_tunnel_with_password(ssh_mock):
    ssh_resource = SSHResource(
        remote_host="remote_host",
        remote_port=12345,
        username="username",
        password="password",
        timeout=10,
        key_file="fake.file",
        keepalive_interval=30,
        compress=True,
        no_host_key_check=False,
        allow_host_key_change=False,
        logger=logging.root.getChild("test_resources"),
    )

    with ssh_resource.get_tunnel(1234):
        ssh_mock.assert_called_once_with(
            "remote_host",
            ssh_port=12345,
            ssh_username="username",
            ssh_password="password",
            ssh_pkey="fake.file",
            ssh_proxy=None,
            local_bind_address=("localhost",),
            remote_bind_address=("localhost", 1234),
            logger=ssh_resource.log,
        )


@mock.patch("dagster_ssh.resources.SSHTunnelForwarder")
def test_tunnel_without_password(ssh_mock):
    ssh_resource = SSHResource(
        remote_host="remote_host",
        remote_port=12345,
        username="username",
        password=None,
        timeout=10,
        key_file="fake.file",
        keepalive_interval=30,
        compress=True,
        no_host_key_check=False,
        allow_host_key_change=False,
        logger=logging.root.getChild("test_resources"),
    )

    with ssh_resource.get_tunnel(1234):
        ssh_mock.assert_called_once_with(
            "remote_host",
            ssh_port=12345,
            ssh_username="username",
            ssh_pkey="fake.file",
            ssh_proxy=None,
            local_bind_address=("localhost",),
            remote_bind_address=("localhost", 1234),
            host_pkey_directories=[],
            logger=ssh_resource.log,
        )


@mock.patch("dagster_ssh.resources.SSHTunnelForwarder")
def test_tunnel_with_string_key(ssh_mock):
    ssh_key = generate_ssh_key()

    ssh_resource = SSHResource(
        remote_host="remote_host",
        remote_port=12345,
        username="username",
        password=None,
        timeout=10,
        key_string=six.ensure_str(ssh_key),
        keepalive_interval=30,
        compress=True,
        no_host_key_check=False,
        allow_host_key_change=False,
        logger=logging.root.getChild("test_resources"),
    )

    with ssh_resource.get_tunnel(1234):
        ssh_mock.assert_called_once_with(
            "remote_host",
            ssh_port=12345,
            ssh_username="username",
            ssh_pkey=key_from_str(ssh_key),
            ssh_proxy=None,
            local_bind_address=("localhost",),
            remote_bind_address=("localhost", 1234),
            host_pkey_directories=[],
            logger=ssh_resource.log,
        )


def test_ssh_sftp(sftpserver):
    tmp_path = get_system_temp_directory()
    readme_file = os.path.join(tmp_path, "readme.txt")

    @solid(
        config_schema={
            "local_filepath": Field(str, is_required=True, description="local file path to get"),
            "remote_filepath": Field(str, is_required=True, description="remote file path to get"),
        },
        required_resource_keys={"ssh_resource"},
    )
    def sftp_solid_get(context):
        local_filepath = context.solid_config.get("local_filepath")
        remote_filepath = context.solid_config.get("remote_filepath")
        return context.resources.ssh_resource.sftp_get(remote_filepath, local_filepath)

    with sftpserver.serve_content({"a_dir": {"readme.txt": "hello, world"}}):
        result = execute_solid(
            sftp_solid_get,
            ModeDefinition(resource_defs={"ssh_resource": sshresource}),
            run_config={
                "solids": {
                    "sftp_solid_get": {
                        "config": {
                            "local_filepath": readme_file,
                            "remote_filepath": "a_dir/readme.txt",
                        }
                    }
                },
                "resources": {
                    "ssh_resource": {
                        "config": {
                            "remote_host": sftpserver.host,
                            "remote_port": sftpserver.port,
                            "username": "user",
                            "password": "pw",
                            "no_host_key_check": True,
                        }
                    }
                },
            },
        )
        assert result.success

    with open(readme_file, "rb") as f:
        contents = f.read()
        assert b"hello, world" in contents
