From 776e625413fe7937853e1c812773f123b0bad9fc Mon Sep 17 00:00:00 2001
From: ikoz <john@kozyrakis.gr>
Date: Tue, 15 Mar 2016 14:58:38 +0000
Subject: [PATCH] Add tests for add-server-certs-to-client-chain feature

---
 test/mitmproxy/test_server.py | 60 +++++++++++++++++++++++++++++++++++
 test/mitmproxy/tservers.py    |  2 ++
 2 files changed, 62 insertions(+)

diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index d7b23bbb8..3286df892 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -999,3 +999,63 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
         # (both terminated)
         # nothing happened here
         assert self.chain[1].tmaster.state.flow_count() == 2
+
+
+class TestHTTPSAddServerCertsToClientChainTrue(tservers.HTTPProxyTest):
+    ssl = True
+    add_server_certs_to_client_chain = True
+    servercert = tutils.test_data.path("data/trusted-server.crt")
+    ssloptions = pathod.SSLOptions(
+            cn="trusted-cert",
+            certs=[
+                ("trusted-cert", servercert)
+            ]
+    )
+
+    def test_add_server_certs_to_client_chain_true(self):
+        """
+        If --add-server-certs-to-client-chain is True, then the client should receive the server's certificates
+        """
+        with open(self.servercert, "rb") as f:
+            d = f.read()
+        c1 = SSLCert.from_pem(d)
+        p = self.pathoc()
+        print("digest of p.cert[1]: %s"%p.server_certs[1].digest('sha256'))
+        print("digest of c1.cert[1]: %s"%c1.digest('sha256'))
+        server_cert_found_in_client_chain = False
+
+        for cert in p.server_certs:
+            if cert.digest('sha256') == c1.digest('sha256'):
+                server_cert_found_in_client_chain = True
+                break
+
+        assert(server_cert_found_in_client_chain == True)
+
+
+class TestHTTPSAddServerCertsToClientChainFalse(tservers.HTTPProxyTest):
+    ssl = True
+    add_server_certs_to_client_chain = False
+    servercert = tutils.test_data.path("data/trusted-server.crt")
+    ssloptions = pathod.SSLOptions(
+            cn="trusted-cert",
+            certs=[
+                ("trusted-cert", servercert)
+            ]
+    )
+
+    def test_add_server_certs_to_client_chain_false(self):
+        """
+        If --add-server-certs-to-client-chain is False, then the client should not receive the server's certificates
+        """
+        with open(self.servercert, "rb") as f:
+            d = f.read()
+        c1 = SSLCert.from_pem(d)
+        p = self.pathoc()
+        server_cert_found_in_client_chain = False
+
+        for cert in p.server_certs:
+            if cert.digest('sha256') == c1.digest('sha256'):
+                server_cert_found_in_client_chain = True
+                break
+
+        assert(server_cert_found_in_client_chain == False)
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index b7b5de9e8..cabd8e1f9 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -86,6 +86,7 @@ class ProxyTestBase(object):
     no_upstream_cert = False
     authenticator = None
     masterclass = TestMaster
+    add_server_certs_to_client_chain = False
 
     @classmethod
     def setup_class(cls):
@@ -129,6 +130,7 @@ class ProxyTestBase(object):
             no_upstream_cert = cls.no_upstream_cert,
             cadir = cls.cadir,
             authenticator = cls.authenticator,
+            add_server_certs_to_client_chain = cls.add_server_certs_to_client_chain,
         )