#!/usr/bin/env python

"""
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 "pylint --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 httplib
import os
import platform
import shutil
import socket
import ssl
import subprocess
import tempfile
import time
import unittest


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

DEVNULL = open("/dev/null", "w")

TEMPDIR = "/does/not/exist"


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

    def gen_certs(self, self_sign=True):
        subprocess.check_call(
            ["openssl", "genrsa", "-out", "%s/key.pem" % self.path, "2048"],
            stdout=DEVNULL, stderr=DEVNULL)

        req_args = ["openssl", "req", "-new", "-batch",
                    "-subj", ("/commonName=*" +
                              "/organizationalUnitName=kxd-tests-%s:%s@%s" % (
                                  self.name, os.getlogin(), platform.node())),
                    "-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])

        subprocess.check_call(req_args, stdout=DEVNULL, stderr=DEVNULL)

    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 open(self.path + "/cert.pem").read()


class CA(object):
    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")
                open("kxd-ca/serial", "w").write("1000\n")
                subprocess.check_output(
                    ["openssl", "req", "-new", "-x509", "-batch",
                     "-config", OPENSSL_CONF,
                     "-subj", ("/commonName=*" +
                               "/organizationalUnitName=kxd-tests-ca:%s@%s" % (
                                   os.getlogin(), platform.node())),
                     "-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 open(self.path + "/cacert.pem").read()


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)
        open(key_path + "key", "w").write(self.keys[name])

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

        if allowed_hosts is not None:
            hfd = open(key_path + "/allowed_hosts", "a")
            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]
    print "Launching server: ", " ".join(args)
    return subprocess.Popen(args)


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

    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:
                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.assertRegexpMatches(err.output, 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.assertEquals(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.assertEquals(key, self.server.keys["k1"])

        key = self.client2.call(self.server.cert_path(), "kxd://localhost/k1")
        self.assertEquals(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.assertEquals(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.assertEquals(data, self.server.keys[key])

            data = self.client2.call(self.server.cert_path(),
                                     "kxd://localhost/%s" % key)
            self.assertEquals(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(open(server1.cert_path()).read())
        server_certs.write(open(server2.cert_path()).read())
        server_certs.close()

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

        self.daemon.kill()
        time.sleep(0.5)
        self.launch_server(server2)

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


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

    def test_tricky(self):
        # No local certificate.
        conn = httplib.HTTPSConnection("localhost", 19840)
        try:
            conn.request("GET", "/v1/")
        except ssl.SSLError as err:
            self.assertRegexpMatches(str(err), "alert bad certificate")
        else:
            self.fail("Client call did not fail as expected")

        # Requests with '..'.
        conn = httplib.HTTPSConnection("localhost", 19840,
                                       key_file=self.client.key_path(),
                                       cert_file=self.client.cert_path())
        conn.request("GET", "/v1/a/../b")
        response = conn.getresponse()

        # Go's http server intercepts these and gives us a 301 Moved
        # Permanently.
        self.assertEquals(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.assertEquals(server_cert, self.server.cert())


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.
        cfd = open(self.server.path + "/data/k1/allowed_clients", "r+")
        for _ in range(5):
            cfd.readline()
        cfd.write('+/+BROKEN+/+')
        cfd.close()

        self.assertClientFails(
            "kxd://localhost/k1",
            "500 Internal Server Error.*Error loading certs")

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

    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.assertEquals(key, self.server.keys["k1"])

        # The server is signed by the CA, but the CA is unknown to the client,
        # so it can't validate it, even if it knows the server directly.
        self.assertClientFails("kxd://localhost/k1",
                               "certificate signed by unknown authority",
                               cert_path=self.server.cert_path())

        # 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.assertEquals(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,
        # so it can't validate it, even if it knows the client directly.
        self.server.new_key("k3",
                            allowed_clients=[self.client.cert()],
                            allowed_hosts=["localhost"])
        self.assertClientFails("kxd://localhost/k3",
                               "403 Forbidden.*No allowed certificate found",
                               cert_path=self.server.cert_path())

    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.assertEquals(key, self.server.keys["k1"])


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