from __future__ import print_function, division, absolute_import

from contextlib import contextmanager
import sys

try:
    import ssl
except ImportError:
    ssl = None

import pytest
from tornado import gen

from distributed.comm import connect, listen
from distributed.security import Security
from distributed.utils_test import new_config, get_cert, gen_test


ca_file = get_cert("tls-ca-cert.pem")

cert1 = get_cert("tls-cert.pem")
key1 = get_cert("tls-key.pem")
keycert1 = get_cert("tls-key-cert.pem")

# Note this cipher uses RSA auth as this matches our test certs
FORCED_CIPHER = "ECDHE-RSA-AES128-GCM-SHA256"

TLS_13_CIPHERS = [
    "TLS_AES_128_GCM_SHA256",
    "TLS_AES_256_GCM_SHA384",
    "TLS_CHACHA20_POLY1305_SHA256",
    "TLS_AES_128_CCM_SHA256",
    "TLS_AES_128_CCM_8_SHA256",
]


def test_defaults():
    with new_config({}):
        sec = Security()
    assert sec.require_encryption in (None, False)
    assert sec.tls_ca_file is None
    assert sec.tls_ciphers is None
    assert sec.tls_client_key is None
    assert sec.tls_client_cert is None
    assert sec.tls_scheduler_key is None
    assert sec.tls_scheduler_cert is None
    assert sec.tls_worker_key is None
    assert sec.tls_worker_cert is None


def test_attribute_error():
    sec = Security()
    assert hasattr(sec, "tls_ca_file")
    with pytest.raises(AttributeError):
        sec.tls_foobar
    with pytest.raises(AttributeError):
        sec.tls_foobar = ""


def test_from_config():
    c = {
        "tls": {
            "ca-file": "ca.pem",
            "scheduler": {"key": "skey.pem", "cert": "scert.pem"},
            "worker": {"cert": "wcert.pem"},
            "ciphers": FORCED_CIPHER,
        },
        "require-encryption": True,
    }
    with new_config(c):
        sec = Security()
    assert sec.require_encryption is True
    assert sec.tls_ca_file == "ca.pem"
    assert sec.tls_ciphers == FORCED_CIPHER
    assert sec.tls_client_key is None
    assert sec.tls_client_cert is None
    assert sec.tls_scheduler_key == "skey.pem"
    assert sec.tls_scheduler_cert == "scert.pem"
    assert sec.tls_worker_key is None
    assert sec.tls_worker_cert == "wcert.pem"


def test_kwargs():
    c = {
        "tls": {
            "ca-file": "ca.pem",
            "scheduler": {"key": "skey.pem", "cert": "scert.pem"},
        }
    }
    with new_config(c):
        sec = Security(
            tls_scheduler_cert="newcert.pem", require_encryption=True, tls_ca_file=None
        )
    assert sec.require_encryption is True
    # None value didn't override default
    assert sec.tls_ca_file == "ca.pem"
    assert sec.tls_ciphers is None
    assert sec.tls_client_key is None
    assert sec.tls_client_cert is None
    assert sec.tls_scheduler_key == "skey.pem"
    assert sec.tls_scheduler_cert == "newcert.pem"
    assert sec.tls_worker_key is None
    assert sec.tls_worker_cert is None


def test_repr():
    with new_config({}):
        sec = Security(tls_ca_file="ca.pem", tls_scheduler_cert="scert.pem")
        assert (
            repr(sec)
            == "Security(tls_ca_file='ca.pem', tls_scheduler_cert='scert.pem')"
        )


def test_tls_config_for_role():
    c = {
        "tls": {
            "ca-file": "ca.pem",
            "scheduler": {"key": "skey.pem", "cert": "scert.pem"},
            "worker": {"cert": "wcert.pem"},
            "ciphers": FORCED_CIPHER,
        }
    }
    with new_config(c):
        sec = Security()
    t = sec.get_tls_config_for_role("scheduler")
    assert t == {
        "ca_file": "ca.pem",
        "key": "skey.pem",
        "cert": "scert.pem",
        "ciphers": FORCED_CIPHER,
    }
    t = sec.get_tls_config_for_role("worker")
    assert t == {
        "ca_file": "ca.pem",
        "key": None,
        "cert": "wcert.pem",
        "ciphers": FORCED_CIPHER,
    }
    t = sec.get_tls_config_for_role("client")
    assert t == {
        "ca_file": "ca.pem",
        "key": None,
        "cert": None,
        "ciphers": FORCED_CIPHER,
    }
    with pytest.raises(ValueError):
        sec.get_tls_config_for_role("supervisor")


def test_connection_args():
    def basic_checks(ctx):
        assert ctx.verify_mode == ssl.CERT_REQUIRED
        assert ctx.check_hostname is False

    def many_ciphers(ctx):
        if sys.version_info >= (3, 6):
            assert len(ctx.get_ciphers()) > 2  # Most likely

    c = {
        "tls": {
            "ca-file": ca_file,
            "scheduler": {"key": key1, "cert": cert1},
            "worker": {"cert": keycert1},
        }
    }
    with new_config(c):
        sec = Security()

    d = sec.get_connection_args("scheduler")
    assert not d["require_encryption"]
    ctx = d["ssl_context"]
    basic_checks(ctx)
    many_ciphers(ctx)

    d = sec.get_connection_args("worker")
    ctx = d["ssl_context"]
    basic_checks(ctx)
    many_ciphers(ctx)

    # No cert defined => no TLS
    d = sec.get_connection_args("client")
    assert d.get("ssl_context") is None

    # With more settings
    c["tls"]["ciphers"] = FORCED_CIPHER
    c["require-encryption"] = True

    with new_config(c):
        sec = Security()

    d = sec.get_listen_args("scheduler")
    assert d["require_encryption"]
    ctx = d["ssl_context"]
    basic_checks(ctx)
    if sys.version_info >= (3, 6):
        supported_ciphers = ctx.get_ciphers()
        tls_12_ciphers = [c for c in supported_ciphers if c["protocol"] == "TLSv1.2"]
        assert len(tls_12_ciphers) == 1
        tls_13_ciphers = [c for c in supported_ciphers if c["protocol"] == "TLSv1.3"]
        if len(tls_13_ciphers):
            assert len(tls_13_ciphers) == 3


def test_listen_args():
    def basic_checks(ctx):
        assert ctx.verify_mode == ssl.CERT_REQUIRED
        assert ctx.check_hostname is False

    def many_ciphers(ctx):
        if sys.version_info >= (3, 6):
            assert len(ctx.get_ciphers()) > 2  # Most likely

    c = {
        "tls": {
            "ca-file": ca_file,
            "scheduler": {"key": key1, "cert": cert1},
            "worker": {"cert": keycert1},
        }
    }
    with new_config(c):
        sec = Security()

    d = sec.get_listen_args("scheduler")
    assert not d["require_encryption"]
    ctx = d["ssl_context"]
    basic_checks(ctx)
    many_ciphers(ctx)

    d = sec.get_listen_args("worker")
    ctx = d["ssl_context"]
    basic_checks(ctx)
    many_ciphers(ctx)

    # No cert defined => no TLS
    d = sec.get_listen_args("client")
    assert d.get("ssl_context") is None

    # With more settings
    c["tls"]["ciphers"] = FORCED_CIPHER
    c["require-encryption"] = True

    with new_config(c):
        sec = Security()

    d = sec.get_listen_args("scheduler")
    assert d["require_encryption"]
    ctx = d["ssl_context"]
    basic_checks(ctx)
    if sys.version_info >= (3, 6):
        supported_ciphers = ctx.get_ciphers()
        tls_12_ciphers = [c for c in supported_ciphers if c["protocol"] == "TLSv1.2"]
        assert len(tls_12_ciphers) == 1
        tls_13_ciphers = [c for c in supported_ciphers if c["protocol"] == "TLSv1.3"]
        if len(tls_13_ciphers):
            assert len(tls_13_ciphers) == 3


@gen_test()
def test_tls_listen_connect():
    """
    Functional test for TLS connection args.
    """

    @gen.coroutine
    def handle_comm(comm):
        peer_addr = comm.peer_address
        assert peer_addr.startswith("tls://")
        yield comm.write("hello")
        yield comm.close()

    c = {
        "tls": {
            "ca-file": ca_file,
            "scheduler": {"key": key1, "cert": cert1},
            "worker": {"cert": keycert1},
        }
    }
    with new_config(c):
        sec = Security()

    c["tls"]["ciphers"] = FORCED_CIPHER
    with new_config(c):
        forced_cipher_sec = Security()

    with listen(
        "tls://", handle_comm, connection_args=sec.get_listen_args("scheduler")
    ) as listener:
        comm = yield connect(
            listener.contact_address, connection_args=sec.get_connection_args("worker")
        )
        msg = yield comm.read()
        assert msg == "hello"
        comm.abort()

        # No SSL context for client
        with pytest.raises(TypeError):
            yield connect(
                listener.contact_address,
                connection_args=sec.get_connection_args("client"),
            )

        # Check forced cipher
        comm = yield connect(
            listener.contact_address,
            connection_args=forced_cipher_sec.get_connection_args("worker"),
        )
        cipher, _, _, = comm.extra_info["cipher"]
        assert cipher in [FORCED_CIPHER] + TLS_13_CIPHERS
        comm.abort()


@gen_test()
def test_require_encryption():
    """
    Functional test for "require_encryption" setting.
    """

    @gen.coroutine
    def handle_comm(comm):
        comm.abort()

    c = {
        "tls": {
            "ca-file": ca_file,
            "scheduler": {"key": key1, "cert": cert1},
            "worker": {"cert": keycert1},
        }
    }
    with new_config(c):
        sec = Security()
    c["require-encryption"] = True
    with new_config(c):
        sec2 = Security()

    for listen_addr in ["inproc://", "tls://"]:
        with listen(
            listen_addr, handle_comm, connection_args=sec.get_listen_args("scheduler")
        ) as listener:
            comm = yield connect(
                listener.contact_address,
                connection_args=sec2.get_connection_args("worker"),
            )
            comm.abort()

        with listen(
            listen_addr, handle_comm, connection_args=sec2.get_listen_args("scheduler")
        ) as listener:
            comm = yield connect(
                listener.contact_address,
                connection_args=sec2.get_connection_args("worker"),
            )
            comm.abort()

    @contextmanager
    def check_encryption_error():
        with pytest.raises(RuntimeError) as excinfo:
            yield
        assert "encryption required" in str(excinfo.value)

    for listen_addr in ["tcp://"]:
        with listen(
            listen_addr, handle_comm, connection_args=sec.get_listen_args("scheduler")
        ) as listener:
            comm = yield connect(
                listener.contact_address,
                connection_args=sec.get_connection_args("worker"),
            )
            comm.abort()

            with pytest.raises(RuntimeError):
                yield connect(
                    listener.contact_address,
                    connection_args=sec2.get_connection_args("worker"),
                )

        with pytest.raises(RuntimeError):
            listen(
                listen_addr,
                handle_comm,
                connection_args=sec2.get_listen_args("scheduler"),
            )