# tests.test_context - test ssl context creation
#
# Copyright 2012 Canonical Ltd.
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License version 3,
# as published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranties of
# MERCHANTABILITY, SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR
# PURPOSE.  See the GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# In addition, as a special exception, the copyright holders give
# permission to link the code of portions of this program with the
# OpenSSL library under certain conditions as described in each
# individual source file, and distribute linked combinations
# including the two.
# You must obey the GNU General Public License in all respects
# for all of the code used other than OpenSSL.  If you modify
# file(s) with this exception, you may extend this exception to your
# version of the file(s), but you are not obligated to do so.  If you
# do not wish to do so, delete this exception statement from your
# version.  If you delete this exception statement from all source
# files in the program, then also delete it here.

import os
import sys

from OpenSSL import crypto, SSL
from twisted.internet import defer, error, reactor, ssl
from twisted.trial import unittest
from twisted.web import client, resource, server

from ubuntuone.storageprotocol import context


class FakeCerts(object):
    """CA and Server certificate."""

    def __init__(self, testcase, common_name="fake.domain"):
        """Initialize this fake instance."""
        self.cert_dir = os.path.join(testcase.mktemp(), 'certs')
        if not os.path.exists(self.cert_dir):
            os.makedirs(self.cert_dir)

        ca_key = self._build_key()
        ca_req = self._build_request(ca_key, "Fake Cert Authority")
        self.ca_cert = self._build_cert(ca_req, ca_req, ca_key)

        server_key = self._build_key()
        server_req = self._build_request(server_key, common_name)
        server_cert = self._build_cert(server_req, self.ca_cert, ca_key)

        self.server_key_path = self._save_key(server_key, "server_key.pem")
        self.server_cert_path = self._save_cert(server_cert, "server_cert.pem")

    def _save_key(self, key, filename):
        """Save a certificate."""
        data = crypto.dump_privatekey(crypto.FILETYPE_PEM, key)
        return self._save(filename, data)

    def _save_cert(self, cert, filename):
        """Save a certificate."""
        data = crypto.dump_certificate(crypto.FILETYPE_PEM, cert)
        return self._save(filename, data)

    def _save(self, filename, data):
        """Save a key or certificate, and return the full path."""
        fullpath = os.path.join(self.cert_dir, filename)
        if os.path.exists(fullpath):
            os.unlink(fullpath)
        with open(fullpath, 'wt') as fd:
            fd.write(data)
        return fullpath

    def _build_key(self):
        """Create a private/public key, save it in a temp dir."""
        key = crypto.PKey()
        key.generate_key(crypto.TYPE_RSA, 1024)
        return key

    def _build_request(self, key, common_name):
        """Create a new certificate request."""
        request = crypto.X509Req()
        request.get_subject().CN = common_name
        request.set_pubkey(key)
        request.sign(key, "md5")
        return request

    def _build_cert(self, request, ca_cert, ca_key):
        """Create a new certificate."""
        certificate = crypto.X509()
        certificate.set_serial_number(1)
        certificate.set_issuer(ca_cert.get_subject())
        certificate.set_subject(request.get_subject())
        certificate.set_pubkey(request.get_pubkey())
        certificate.gmtime_adj_notBefore(0)
        certificate.gmtime_adj_notAfter(3600)  # valid for one hour
        certificate.sign(ca_key, "md5")
        return certificate


class FakeResource(resource.Resource):
    """A fake resource."""

    isLeaf = True

    def render(self, request):
        """Render this resource."""
        return "ok"


class SSLContextTestCase(unittest.TestCase):
    """Tests for the context.get_ssl_context function."""

    @defer.inlineCallbacks
    def setUp(self):
        yield super(SSLContextTestCase, self).setUp()
        self.patch(context, "get_certificates", self.get_certificates)

    def get_certificates(self):
        """Get the uninstalled certificates, for testing."""
        certs = []
        data_dir = os.path.join(os.path.dirname(__file__), os.pardir, 'data')
        for certname in ['UbuntuOne-Go_Daddy_Class_2_CA.pem',
                         'UbuntuOne-Go_Daddy_CA.pem']:
            cert_path = os.path.abspath(os.path.join(data_dir, certname))
            c = ssl.Certificate.loadPEM(file(cert_path, 'r').read())
            certs.append(c.original)

        return certs

    @defer.inlineCallbacks
    def verify_context(self, server_context, client_context):
        """Verify a client context with a given server context."""
        site = server.Site(FakeResource())
        port = reactor.listenSSL(0, site, server_context)
        self.addCleanup(port.stopListening)
        url = "https://localhost:%d" % port.getHost().port
        result = yield client.getPage(url, contextFactory=client_context)
        self.assertEqual(result, "ok")

    @defer.inlineCallbacks
    def test_no_verify(self):
        """Test the no_verify option."""
        certs = FakeCerts(self, "localhost")
        server_context = ssl.DefaultOpenSSLContextFactory(
            certs.server_key_path, certs.server_cert_path)
        client_context = context.get_ssl_context(no_verify=True,
                                                 hostname="localhost")

        yield self.verify_context(server_context, client_context)

    def test_no_hostname(self):
        """Test that calling without hostname arg raises proper error."""
        self.assertRaises(error.CertificateError,
                          context.get_ssl_context, False)

    @defer.inlineCallbacks
    def test_fails_certificate(self):
        """A wrong certificate is rejected."""
        certs = FakeCerts(self, "localhost")
        server_context = ssl.DefaultOpenSSLContextFactory(
            certs.server_key_path, certs.server_cert_path)
        client_context = context.get_ssl_context(no_verify=False,
                                                 hostname="localhost")

        d = self.verify_context(server_context, client_context)
        e = yield self.assertFailure(d, SSL.Error)
        self.assertEqual(e[0][0][1], "SSL3_GET_SERVER_CERTIFICATE")

    @defer.inlineCallbacks
    def test_fails_hostname(self):
        """A wrong hostname is rejected."""
        certs = FakeCerts(self, "thisiswronghost.net")
        server_context = ssl.DefaultOpenSSLContextFactory(
            certs.server_key_path, certs.server_cert_path)
        self.patch(context, "get_certificates", lambda: [certs.ca_cert])
        client_context = context.get_ssl_context(no_verify=False,
                                                 hostname="localhost")

        d = self.verify_context(server_context, client_context)
        e = yield self.assertFailure(d, SSL.Error)
        self.assertEqual(e[0][0][1], "SSL3_GET_SERVER_CERTIFICATE")

    @defer.inlineCallbacks
    def test_matches_all(self):
        """A valid certificate passes checks."""
        certs = FakeCerts(self, "localhost")
        server_context = ssl.DefaultOpenSSLContextFactory(
            certs.server_key_path, certs.server_cert_path)
        self.patch(context, "get_certificates", lambda: [certs.ca_cert])
        client_context = context.get_ssl_context(no_verify=False,
                                                 hostname="localhost")

        yield self.verify_context(server_context, client_context)


class SSLCertLocationTestCase(unittest.TestCase):
    """Test determining the cert location."""

    @defer.inlineCallbacks
    def setUp(self):
        yield super(SSLCertLocationTestCase, self).setUp()

    def test_win(self):
        """Test geting a path when Common AppData is defined."""
        self.patch(context, "load_config_paths",
                   lambda x: [os.path.join("returned-value", x)])
        self.patch(sys, "platform", "win32")
        path = context.get_cert_location()
        self.assertEqual(path, os.path.join("returned-value",
                                            "ubuntuone-storageprotocol"))

    def test_darwin_frozen(self):
        """Test that we get a path with .app in it on frozen darwin."""
        self.patch(sys, "platform", "darwin")
        sys.frozen = "macosx-app"
        self.addCleanup(delattr, sys, "frozen")
        self.patch(context, "__file__",
                   os.path.join("path", "to", "Main.app", "ignore"))
        path = context.get_cert_location()
        self.assertEqual(path, os.path.join("path", "to", "Main.app",
                                            "Contents", "Resources"))

    def test_darwin_unfrozen(self):
        """Test that we get a source-relative path on unfrozen darwin."""
        self.patch(sys, "platform", "darwin")
        self.patch(context, "__file__",
                   os.path.join("path", "to", "ubuntuone",
                                "storageprotocol", "context.py"))
        path = context.get_cert_location()
        self.assertEqual(path, os.path.join("path", "to", "data"))

    def test_linux(self):
        """Test that linux gets the right path."""
        self.patch(sys, "platform", "linux2")
        path = context.get_cert_location()
        self.assertEqual(path, "/etc/ssl/certs")
