#!/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 "pylint3 --rcfile=.pylintrc run_tests" after making changes,
# to make sure the file has a reasonably uniform coding style.
# You can also use "autopep8 -d --ignore=E301,E26 run_tests" to help with
# this, but make sure the output looks sane.


import contextlib
import http.client
import os
import shutil
import socket
import ssl
import subprocess
import tempfile
import textwrap
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")

# Path to the test OpenSSL configuration.
OPENSSL_CONF = os.path.abspath(
    os.path.dirname(os.path.realpath(__file__)) + "/openssl.cnf")

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

    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):
        self.path = tempfile.mkdtemp(prefix="config-%s-" % name, dir=TEMPDIR)
        self.name = name

    def gen_certs(self, self_sign=True):
        try:
            cmd = ["openssl", "genrsa",
                   "-out", "%s/key.pem" % self.path, "2048"]
            subprocess.check_output(cmd, stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as err:
            print("openssl call failed, output: %r" % err.output)
            raise

        ouname = "kxd-tests-%s" % self.name
        req_args = ["openssl", "req", "-new", "-batch",
                    "-subj", ("/commonName=*" +
                              "/organizationalUnitName=%s" % ouname),
                    "-key", "%s/key.pem" % self.path]
        if self_sign:
            req_args.extend(["-x509", "-out", "%s/cert.pem" % self.path])
        else:
            req_args.extend(["-out", "%s/cert.csr" % self.path])

        try:
            subprocess.check_output(req_args, stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as err:
            print("openssl 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 csr_path(self):
        return self.path + "/cert.csr"

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


class CA:
    def __init__(self):
        self.path = tempfile.mkdtemp(prefix="config-ca-", dir=TEMPDIR)
        os.makedirs(self.path + "/kxd-ca/newcerts/")

        try:
            # We need to run the CA commands from within the path.
            with pushd(self.path):
                open("kxd-ca/index.txt", "w").close()
                with open("kxd-ca/serial", "w") as serial:
                    serial.write("1000\n")
                subprocess.check_output(
                    ["openssl", "req", "-new", "-x509", "-batch",
                     "-config", OPENSSL_CONF,
                     "-subj", ("/commonName=*" +
                               "/organizationalUnitName=kxd-tests-ca"),
                     "-extensions", "v3_ca", "-nodes",
                     "-keyout", "cakey.pem",
                     "-out", "cacert.pem"],
                    stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as err:
            print("openssl call failed, output: %r" % err.output)
            raise

    def sign(self, csr):
        try:
            with pushd(self.path):
                subprocess.check_output(
                    ["openssl", "ca", "-batch", "-config", OPENSSL_CONF,
                     "-keyfile", "cakey.pem", "-cert", "cacert.pem",
                     "-in", csr, "-out", "%s.pem" % os.path.splitext(csr)[0]],
                    stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as err:
            print("openssl call failed, output: %r" % err.output)
            raise

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

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


class ServerConfig(Config):
    def __init__(self, self_sign=True, name="server"):
        Config.__init__(self, name)
        self.keys = {}
        self.gen_certs(self_sign)

    def new_key(self, name, allowed_clients=None, allowed_hosts=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")


class ClientConfig(Config):
    def __init__(self, self_sign=True, name="client"):
        Config.__init__(self, name)
        self.gen_certs(self_sign)

    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


def launch_daemon(cfg):
    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]
    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()


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.kill()
            self.daemon.wait()

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

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

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

        # Unknown key -> 404.
        self.assertClientFails("kxd://localhost/k2", "404 Not Found")

        # No certificates allowed -> 403.
        self.server.new_key("k3", allowed_hosts=["localhost"])
        self.assertClientFails("kxd://localhost/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("kxd://localhost/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("kxd://localhost/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("kxd://localhost/k1",
                               "certificate signed by unknown authority",
                               cert_path=self.client.cert_path())


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.kill()
        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.close()
        except ssl.SSLError as err:
            self.assertEqual(err.reason, "SSLV3_ALERT_BAD_CERTIFICATE")
        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_server_cert(self):
        rawsock = socket.create_connection(("localhost", 19840))
        sock = ssl.wrap_socket(rawsock,
                               keyfile=self.client.key_path(),
                               certfile=self.client.cert_path())

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


class Delegation(TestCase):
    """Tests for CA delegations."""
    def setUp(self):
        # For these tests, we don't have a common setup, as each will create
        # server and clients in a slightly different way.
        self.daemon = None

    def prepare(self, server_self_sign=True, client_self_sign=True,
                ca_sign_server=None, ca_sign_client=None):
        self.server = ServerConfig(self_sign=server_self_sign)
        self.client = ClientConfig(self_sign=client_self_sign)

        self.ca = CA()
        if ca_sign_server is None:
            ca_sign_server = not server_self_sign
        if ca_sign_client is None:
            ca_sign_client = not client_self_sign

        if ca_sign_server:
            self.ca.sign(self.server.csr_path())
        if ca_sign_client:
            self.ca.sign(self.client.csr_path())

        self.launch_server(self.server)

    def test_server_delegated(self):
        self.prepare(server_self_sign=False)

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

        # Successful request.
        key = self.client.call(self.ca.cert_path(), "kxd://localhost/k1")
        self.assertEqual(key, self.server.keys["k1"])

        # The server is signed by the CA, but the CA is unknown to the client.
        # But the client knows the server directly, so it's allowed.
        key = self.client.call(self.server.cert_path(), "kxd://localhost/k1")
        self.assertEqual(key, self.server.keys["k1"])

        # Same as above, but give the wrong CA.
        ca2 = CA()
        self.assertClientFails("kxd://localhost/k1",
                               "certificate signed by unknown authority",
                               cert_path=ca2.cert_path())

    def test_client_delegated(self):
        self.prepare(client_self_sign=False)

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

        # The CA signing the client is unknown to the server.
        ca2 = CA()
        self.server.new_key("k2",
                            allowed_clients=[ca2.cert()],
                            allowed_hosts=["localhost"])
        self.assertClientFails("kxd://localhost/k2",
                               "403 Forbidden.*No allowed certificate found",
                               cert_path=self.server.cert_path())

        # The client is signed by the CA, but the CA is unknown to the server.
        # But the server it knows the client directly, so it's allowed.
        self.server.new_key("k3",
                            allowed_clients=[self.client.cert()],
                            allowed_hosts=["localhost"])
        key = self.client.call(self.server.cert_path(), "kxd://localhost/k3")
        self.assertEqual(key, self.server.keys["k3"])

    def test_both_delegated(self):
        self.prepare(server_self_sign=False, client_self_sign=False)
        self.server.new_key("k1",
                            allowed_clients=[self.ca.cert()],
                            allowed_hosts=["localhost"])

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


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=OU=kxd-tests-client", 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")


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