#!/usr/bin/env python3

"""
Tests for kxd and kxc
---------------------

This file contains various integration and validation tests for kxc and kxd.

It will create different test configurations and run the compiled server and
client under various conditions, to make sure they behave as intended.
"""

# NOTE: Please run "black run_tests" after making changes, to to make sure the
# file has a reasonably uniform coding style.


import contextlib
import http.client
import os
import shutil
import socket
import ssl
import subprocess
import tempfile
import textwrap
import threading
import time
import tracemalloc
import unittest

tracemalloc.start()

############################################################
# Test infrastructure.
#
# These functions and classes are used to make the individual tests easier to
# write.  For the individual test cases, see below.

# Path to our built binaries; used to run the server and client for testing
# purposes.
BINS = os.path.abspath(os.path.dirname(os.path.realpath(__file__)) + "/../out")

TEMPDIR = "/does/not/exist"

# User the script is running as. Just informational, for troubleshooting
# purposes, so we don't care if it's missing.
LOGNAME = os.environ.get("LOGNAME", "unknown")


def setUpModule():  # pylint: disable=invalid-name
    if not os.path.isfile(BINS + "/kxd"):
        raise RuntimeError("kxd not found at " + BINS + "/kxd")
    if not os.path.isfile(BINS + "/kxc"):
        raise RuntimeError("kxc not found at " + BINS + "/kxc")
    if not os.path.isfile(BINS + "/kxgencert"):
        raise RuntimeError("kxgencert not found at " + BINS + "/kxgencert")

    global TEMPDIR  # pylint: disable=global-statement
    TEMPDIR = tempfile.mkdtemp(prefix="kxdtest-")


def tearDownModule():  # pylint: disable=invalid-name
    # Remove the temporary directory only on success.
    # Be extra paranoid about removing.
    # TODO: Only remove on success.
    if os.environ.get("KEEPTMP"):
        return
    if len(TEMPDIR) > 10 and not TEMPDIR.startswith("/home"):
        shutil.rmtree(TEMPDIR)


@contextlib.contextmanager
def pushd(path):
    prev = os.getcwd()
    os.chdir(path)
    yield
    os.chdir(prev)


class Config:
    def __init__(self, name, host=""):
        self.path = tempfile.mkdtemp(prefix="config-%s-" % name, dir=TEMPDIR)
        self.name = name
        self.host = host
        self.keys = {}

    def gen_cert(self):
        try:
            cmd = [
                BINS + "/kxgencert",
                "-organization=kxd-tests-%s" % self.name,
                "-key=" + self.key_path(),
                "-cert=" + self.cert_path(),
            ]
            if self.host:
                cmd.append("-host=" + self.host)
            subprocess.check_output(cmd, stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as err:
            print("kxgencert call failed, output: %r" % err.output)
            raise

    def cert_path(self):
        return self.path + "/cert.pem"

    def key_path(self):
        return self.path + "/key.pem"

    def cert(self):
        return read_all(self.path + "/cert.pem")

    def new_key(
        self, name, allowed_clients=None, allowed_hosts=None, email_to=None
    ):
        self.keys[name] = os.urandom(1024)
        key_path = self.path + "/data/" + name + "/"
        if not os.path.isdir(key_path):
            os.makedirs(key_path)
        with open(key_path + "key", "bw") as key:
            key.write(self.keys[name])

        if allowed_clients is not None:
            with open(key_path + "/allowed_clients", "a") as cfd:
                for cli in allowed_clients:
                    cfd.write(cli)

        if allowed_hosts is not None:
            with open(key_path + "/allowed_hosts", "a") as hfd:
                for host in allowed_hosts:
                    hfd.write(host + "\n")

        if email_to is not None:
            with open(key_path + "/email_to", "a") as efd:
                efd.write(email_to + "\n")

    def call(self, server_cert, url):
        args = [
            BINS + "/kxc",
            "--client_cert=%s/cert.pem" % self.path,
            "--client_key=%s/key.pem" % self.path,
            "--server_cert=%s" % server_cert,
            url,
        ]
        try:
            print("Running client:", " ".join(args))
            return subprocess.check_output(args, stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as err:
            print("Client call failed, output: %r" % err.output)
            raise


class ServerConfig(Config):
    def __init__(self, name="server", host="localhost"):
        Config.__init__(self, name, host)
        self.gen_cert()

    def call(self, server_cert, url):
        # To prevent accidental calls and to enforce that the test cases
        # don't mix server and client.
        return NotImplementedError("ServerConfig does not support call")


class ClientConfig(Config):
    def __init__(self, name="client", host="localhost"):
        Config.__init__(self, name, host)
        self.gen_cert()

    def new_key(self, name, allowed_clients=None, allowed_hosts=None):
        # To prevent accidental calls and to enforce that the test cases
        # don't mix server and client.
        raise NotImplementedError("ClientConfig does not support new_key")


class StaticConfig(Config):
    def __init__(self, path):
        self.path = path
        self.host = "localhost"
        self.keys = {}

    def gen_cert(self):
        raise NotImplementedError("StaticConfig does not support gen_cert")


def launch_daemon(cfg, smtp_addr=None):
    args = [
        BINS + "/kxd",
        "--data_dir=%s/data" % cfg,
        "--key=%s/key.pem" % cfg,
        "--cert=%s/cert.pem" % cfg,
        "--logfile=%s/log" % cfg,
        "--hook=%s/hook" % cfg,
    ]
    if smtp_addr:
        args.append("--smtp_addr=%s:%s" % smtp_addr)
    print("Launching server: ", " ".join(args))
    return subprocess.Popen(args)


def read_all(fname):
    with open(fname) as fd:  # pylint: disable=invalid-name
        return fd.read()


def receive_emails():
    """Receive emails from the server.

    Return:
    - The server address.
    - A list where we will put the emails: each email is a tuple of
      (destination addresses, body).
    """
    server = socket.create_server(("localhost", 0))
    server.listen(1)
    addr = server.getsockname()

    emails = []

    def handler():
        # Complete a simple SMTP transaction.
        sock, addr = server.accept()
        sock.sendall(b"220 localhost SMTP\n")
        body = b""
        rcpt_to = []
        while True:
            data = sock.recv(1024)
            if not data:
                break
            if data.startswith(b"HELO"):
                sock.sendall(b"250 localhost\n")
            elif data.startswith(b"MAIL FROM"):
                sock.sendall(b"250 OK\n")
            elif data.startswith(b"RCPT TO"):
                rcpt_to.append(data.split(b"<")[1].split(b">")[0].decode())
                sock.sendall(b"250 OK\n")
            elif data.startswith(b"DATA"):
                sock.sendall(b"354 Start data\n")
                body = b""
                while True:
                    body += sock.recv(1024)
                    if body.endswith(b"\r\n.\r\n"):
                        sock.sendall(b"250 OK\n")
                        break
            elif data.startswith(b"QUIT"):
                sock.sendall(b"221 Chau\n")
                break
            else:
                sock.sendall(b"500 Unknown command\n")

        emails.append((rcpt_to, body))
        sock.close()
        server.close()

    server_thread = threading.Thread(target=handler, daemon=True)
    server_thread.start()

    return addr, emails


class TestCase(unittest.TestCase):
    def setUp(self):
        self.server = ServerConfig()
        self.client = ClientConfig()
        self.daemon = None
        self.ca = None  # pylint: disable=invalid-name
        self.launch_server(self.server)

    def tearDown(self):
        if self.daemon:
            self.daemon.terminate()
            self.daemon.wait()

    def launch_server(self, server, smtp_addr=None):
        self.daemon = launch_daemon(server.path, smtp_addr)

        # Wait for the server to start accepting connections.
        deadline = time.time() + 5
        while time.time() < deadline:
            try:
                with socket.create_connection(("localhost", 19840), timeout=5):
                    break
            except socket.error:
                continue
        else:
            self.fail("Timeout waiting for the server")

    # pylint: disable=invalid-name
    def assertClientFails(self, url, regexp, client=None, cert_path=None):
        if client is None:
            client = self.client
        if cert_path is None:
            cert_path = self.server.cert_path()

        try:
            client.call(cert_path, url)
        except subprocess.CalledProcessError as err:
            self.assertRegex(err.output.decode(), regexp)
        else:
            self.fail("Client call did not fail as expected")


############################################################
# Test cases.
#


class Simple(TestCase):
    """Simple test cases for common (mis)configurations."""

    def test_simple(self):
        # There's no need to split these up; by doing all these within a
        # single test, we speed things up significantly, as we avoid the
        # overhead of creating the certificates and bringing up the server.

        server_url = "kxd://" + self.server.host

        # Normal successful case.
        self.server.new_key(
            "k1",
            allowed_clients=[self.client.cert()],
            # Depending on the self.server.host, the client may connect to it
            # from localhost or from another IP; since it is unpredictable we
            # should allow both.
            allowed_hosts=["localhost", self.server.host],
        )
        key = self.client.call(self.server.cert_path(), server_url + "/k1")
        self.assertEqual(key, self.server.keys["k1"])

        # Unknown key -> 404.
        self.assertClientFails(server_url + "/k2", "404 Not Found")

        # No certificates allowed -> 403.
        self.server.new_key("k3", allowed_hosts=["localhost", self.server.host])
        self.assertClientFails(
            server_url + "/k3", "403 Forbidden.*No allowed certificate found"
        )

        # Host not allowed -> 403.
        self.server.new_key(
            "k4", allowed_clients=[self.client.cert()], allowed_hosts=[]
        )
        self.assertClientFails(
            server_url + "/k4", "403 Forbidden.*Host not allowed"
        )

        # Nothing allowed -> 403.
        # We don't restrict the reason of failure, that's not defined in this
        # case, as it could be either the host or the cert that are validated
        # first.
        self.server.new_key("k5")
        self.assertClientFails(server_url + "/k5", "403 Forbidden")

        # We tell the client to expect the server certificate to be the client
        # one, which is never going to work.
        self.assertClientFails(
            server_url + "/k1",
            "certificate signed by unknown authority",
            cert_path=self.client.cert_path(),
        )


class WildcardHostnamesKxgencert(Simple):
    """Tests for certificates with DNSNames='*'.

    In kxd <= 0.16, by default kxgencert (and before than, the equivalent
    scripts) would generate certificates with DNSNames='*', and everything
    worked okay.

    But in Go 1.23, the Go TLS library started to reject such certificates (Go
    commit 375031d8dcec9ae74d2dbc437b201107dba3bb5f).

    We changed the defaults since then, but we still want to make sure that
    the server can handle such certificates, as they might still be in use in
    the wild.

    This test uses static certificates generated with the old defaults, to
    be absolutely sure they survive even through changes in the key generation
    code.
    """

    _cert_dir = "wildcard_test_certs_kxgencert"

    def setUp(self):
        test_dir = os.path.dirname(os.path.realpath(__file__))
        self.server = StaticConfig(test_dir + "/" + self._cert_dir + "/server")
        self.client = StaticConfig(test_dir + "/" + self._cert_dir + "/client")

        self.daemon = None
        self.ca = None  # pylint: disable=invalid-name
        self.launch_server(self.server)

    def test_minimal(self):
        self.server.new_key(
            "k1",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost"],
        )
        key = self.client.call(self.server.cert_path(), "kxd://localhost/k1")
        self.assertEqual(key, self.server.keys["k1"])

    # Reuse the rest of the Simple test cases.


class WildcardHostnamesOpenSSL(WildcardHostnamesKxgencert):
    """Tests for certificates with DNSNames='*' generated with OpenSSL."""

    _cert_dir = "wildcard_test_certs_openssl"


class SpecificHostnames(Simple):
    """Tests for certificates with specific hostnames.

    These tests check that the server ignores the client certificate's
    hostname, and that the client validates the server certificate's hostname
    in scenarios that it is different than `localhost`.
    """

    def setUp(self):
        self.server = ServerConfig(host=socket.gethostname())
        self.client = ClientConfig(host="client-" + os.urandom(8).hex())
        self.daemon = None
        self.ca = None  # pylint: disable=invalid-name
        self.launch_server(self.server)

    def test_minimal(self):
        self.server.new_key(
            "k1",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost", self.server.host],
        )
        key = self.client.call(
            self.server.cert_path(), "kxd://" + self.server.host + "/k1"
        )
        self.assertEqual(key, self.server.keys["k1"])


class CertFor127_0_0_1(Simple):
    """Tests for server certificate for 127.0.0.1 (instead of a name)."""

    def setUp(self):
        self.server = ServerConfig(host="127.0.0.1")
        self.client = ClientConfig(host="client-" + os.urandom(8).hex())
        self.daemon = None
        self.ca = None  # pylint: disable=invalid-name
        self.launch_server(self.server)

    def test_minimal(self):
        self.server.new_key(
            "k1",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost"],
        )
        key = self.client.call(
            self.server.cert_path(), "kxd://" + self.server.host + "/k1"
        )
        self.assertEqual(key, self.server.keys["k1"])


class CommonErrors(TestCase):
    """Simple test cases for common errors."""

    def test_invalid_schema(self):
        self.assertClientFails("invalidschema://a/b", "unsupported URL schema")

    def test_invalid_url(self):
        self.assertClientFails(
            " http://invalid/url",
            "first path segment in URL cannot contain colon",
        )

    def test_empty_server_cert(self):
        self.assertClientFails(
            "kxd://localhost/k1",
            "Failed to load server certs: error appending certificates",
            cert_path="/dev/null",
        )

    def test_unknown_server_cert_file(self):
        self.assertClientFails(
            "kxd://localhost/k1",
            "open /does/not/exist: no such file or directory",
            cert_path="/does/not/exist",
        )

    def test_bad_server_cert_file(self):
        test_dir = os.path.dirname(os.path.realpath(__file__))
        self.assertClientFails(
            "kxd://localhost/k1",
            "x509: malformed certificate",
            cert_path=test_dir + "/bad_cert/cert.pem",
        )


class Multiples(TestCase):
    """Tests for multiple clients and keys."""

    def setUp(self):
        TestCase.setUp(self)
        self.client2 = ClientConfig(name="client2")

    def test_two_clients(self):
        self.server.new_key(
            "k1",
            allowed_clients=[self.client.cert(), self.client2.cert()],
            allowed_hosts=["localhost"],
        )
        key = self.client.call(self.server.cert_path(), "kxd://localhost/k1")
        self.assertEqual(key, self.server.keys["k1"])

        key = self.client2.call(self.server.cert_path(), "kxd://localhost/k1")
        self.assertEqual(key, self.server.keys["k1"])

        # Only one client allowed.
        self.server.new_key(
            "k2",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost"],
        )
        key = self.client.call(self.server.cert_path(), "kxd://localhost/k2")
        self.assertEqual(key, self.server.keys["k2"])

        self.assertClientFails(
            "kxd://localhost/k2",
            "403 Forbidden.*No allowed certificate found",
            client=self.client2,
        )

    def test_many_keys(self):
        keys = ["a", "d/e", "a/b/c", "d/"]
        for key in keys:
            self.server.new_key(
                key,
                allowed_clients=[self.client.cert(), self.client2.cert()],
                allowed_hosts=["localhost"],
            )

        for key in keys:
            data = self.client.call(
                self.server.cert_path(), "kxd://localhost/%s" % key
            )
            self.assertEqual(data, self.server.keys[key])

            data = self.client2.call(
                self.server.cert_path(), "kxd://localhost/%s" % key
            )
            self.assertEqual(data, self.server.keys[key])

        self.assertClientFails("kxd://localhost/a/b", "404 Not Found")

    def test_two_servers(self):
        server1 = self.server
        server1.new_key("k1", allowed_clients=[self.client.cert()])
        server2 = ServerConfig(name="server2")
        server2.new_key("k1", allowed_clients=[self.client.cert()])

        # Write a file containing the certs of both servers.
        server_certs_path = self.client.path + "/server_certs.pem"
        server_certs = open(server_certs_path, "w")
        server_certs.write(read_all(server1.cert_path()))
        server_certs.write(read_all(server2.cert_path()))
        server_certs.close()

        key = self.client.call(server_certs_path, "kxd://localhost/k1")
        self.assertEqual(key, server1.keys["k1"])

        self.daemon.terminate()
        self.daemon.wait()
        time.sleep(0.5)
        self.launch_server(server2)

        key = self.client.call(server_certs_path, "kxd://localhost/k1")
        self.assertEqual(key, server2.keys["k1"])


class TrickyRequests(TestCase):
    """Tests for tricky requests."""

    def https_connection(self, host, port, key_file=None, cert_file=None):
        # Get an SSL context that can validate our server certificate.
        context = ssl.create_default_context(cafile=self.server.cert_path())
        context.check_hostname = False
        if cert_file:
            context.load_cert_chain(cert_file, key_file)
        return http.client.HTTPSConnection(host, port, context=context)

    def test_no_local_cert(self):
        """No local certificate."""
        conn = self.https_connection("localhost", 19840)
        try:
            conn.request("GET", "/v1/")
            conn.getresponse()
            conn.close()
        except ssl.SSLError as err:
            # Expect one of these errors (the specific one can change
            # depending on the version of OpenSSL).
            self.assertIn(
                err.reason,
                [
                    "SSLV3_ALERT_BAD_CERTIFICATE",
                    "TLSV13_ALERT_CERTIFICATE_REQUIRED",
                ],
            )
        else:
            self.fail("Client call did not fail as expected")

    def test_path_with_dotdot(self):
        """Requests with '..'."""
        conn = self.https_connection(
            "localhost",
            19840,
            key_file=self.client.key_path(),
            cert_file=self.client.cert_path(),
        )
        conn.request("GET", "/v1/a/../b")
        response = conn.getresponse()
        conn.close()

        # Go's http server intercepts these and gives us a 301 Moved
        # Permanently.
        self.assertEqual(response.status, 301)

    def test_invalid_path(self):
        conn = self.https_connection(
            "localhost",
            19840,
            key_file=self.client.key_path(),
            cert_file=self.client.cert_path(),
        )
        conn.request("GET", "/v1/a..b")
        response = conn.getresponse()
        conn.close()
        self.assertEqual(response.status, 406)

    def test_server_cert(self):
        rawsock = socket.create_connection(("localhost", 19840))
        context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
        context.load_verify_locations(self.server.cert_path())
        context.check_hostname = False
        sock = context.wrap_socket(rawsock)

        # We don't check the cipher itself, as it depends on the environment,
        # but we should be using >= 128 bit secrets.
        self.assertTrue(sock.cipher()[2] >= 128)

        server_cert = ssl.DER_cert_to_PEM_cert(
            sock.getpeercert(binary_form=True)
        )
        self.assertEqual(server_cert, self.server.cert())
        sock.close()


class BrokenServerConfig(TestCase):
    """Tests for a broken server config."""

    def test_broken_client_certs(self):
        self.server.new_key(
            "k1",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost"],
        )

        # Corrupt the client certificate.
        with open(self.server.path + "/data/k1/allowed_clients", "tr+") as cfd:
            cfd.seek(30)
            cfd.write("+/+BROKEN+/+")

        self.assertClientFails(
            "kxd://localhost/k1",
            "Error loading certs|No allowed certificate found",
        )

    def test_missing_key(self):
        self.server.new_key(
            "k1",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost"],
        )

        os.unlink(self.server.path + "/data/k1/key")
        self.assertClientFails("kxd://localhost/k1", "404 Not Found")


EMAIL_TO_FILE = textwrap.dedent(
    """
    # Comment.
    me@example.com
    you@test.net
    """.strip()
)


class Hook(TestCase):
    """Test cases for hook support."""

    HOOK_SCRIPT_TMPL = textwrap.dedent(
        """
        #!/bin/sh
        pwd > hook-output
        env >> hook-output
        exit {exit_code}
        """.strip()
    )

    def write_hook(self, exit_code):
        path = self.server.path + "/hook"
        script = self.HOOK_SCRIPT_TMPL.format(exit_code=exit_code)

        with open(path, "w") as hook:
            hook.write(script)
        os.chmod(path, 0o770)

    def test_simple(self):
        self.write_hook(exit_code=0)

        # Normal successful case.
        self.server.new_key(
            "k1",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost"],
        )
        key = self.client.call(self.server.cert_path(), "kxd://localhost/k1")
        self.assertEqual(key, self.server.keys["k1"])

        hook_out = read_all(self.server.path + "/data/hook-output")
        self.assertIn("CLIENT_CERT_SUBJECT=O=kxd-tests-client", hook_out)
        self.assertNotIn("EMAIL_TO=", hook_out)

        # Failure caused by the hook exiting with error.
        self.write_hook(exit_code=1)
        self.assertClientFails("kxd://localhost/k1", "Prevented by hook")

        # Failure caused by the hook not being executable.
        self.write_hook(exit_code=0)
        os.chmod(self.server.path + "/hook", 0o660)
        self.assertClientFails("kxd://localhost/k1", "Prevented by hook")

    def test_email_to(self):
        self.write_hook(exit_code=0)
        self.server.new_key(
            "k2",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost"],
            email_to=EMAIL_TO_FILE,
        )
        key = self.client.call(self.server.cert_path(), "kxd://localhost/k2")
        self.assertEqual(key, self.server.keys["k2"])

        hook_out = read_all(self.server.path + "/data/hook-output")
        self.assertIn("CLIENT_CERT_SUBJECT=O=kxd-tests-client", hook_out)
        self.assertIn("EMAIL_TO=me@example.com you@test.net", hook_out)


class Emails(TestCase):
    """Tests for email notifications."""

    def setUp(self):
        self.smtp_addr, self.emails = receive_emails()
        self.server = ServerConfig()
        self.client = ClientConfig()
        self.daemon = None
        self.ca = None  # pylint: disable=invalid-name
        self.launch_server(self.server, smtp_addr=self.smtp_addr)

    def test_emails(self):
        self.server.new_key(
            "k1",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost", self.server.host],
            email_to=EMAIL_TO_FILE,
        )
        key = self.client.call(
            self.server.cert_path(), "kxd://" + self.server.host + "/k1"
        )
        self.assertEqual(key, self.server.keys["k1"])
        self.assertEqual(len(self.emails), 1)
        self.assertEqual(self.emails[0][0], ["me@example.com", "you@test.net"])
        self.assertRegex(
            self.emails[0][1].decode().replace("\r\n", "\n"),
            textwrap.dedent(
                """\
                Date: .*
                From: Key Exchange Daemon <kxd@127.0.0.1>
                To: me@example.com, you@test.net
                Subject: Access to key k1

                Key: k1
                Accessed by: .*
                On: .*

                Client certificate:
                  Signature: .*
                  Subject: O=kxd-tests-client

                Authorizing chains:
                  .*
            """
            ),
        )

    def test_no_emails(self):
        self.server.new_key(
            "k2",
            allowed_clients=[self.client.cert()],
            allowed_hosts=["localhost", self.server.host],
        )
        key = self.client.call(
            self.server.cert_path(), "kxd://" + self.server.host + "/k2"
        )
        self.assertEqual(key, self.server.keys["k2"])
        # Note that we did not set up email_to, so we don't expect any emails.
        self.assertEqual(self.emails, [])


if __name__ == "__main__":
    unittest.main()