From d4f4cfe2252ac82bc361c098e7fba48d70466890 Mon Sep 17 00:00:00 2001
From: Abcdefghijklmnopqrstuvwxyzxyz <huang_yue_zhi@outlook.com>
Date: Fri, 9 Nov 2018 16:06:04 +0800
Subject: [PATCH] Add Organization field for the generated certificate (#3376)

add organization field for the generated certificate
---
 mitmproxy/certs.py              | 20 +++++++++++++++++---
 mitmproxy/proxy/protocol/tls.py |  5 ++++-
 test/mitmproxy/test_certs.py    |  9 +++++++--
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/mitmproxy/certs.py b/mitmproxy/certs.py
index 8b8ba6f2b..6970618e4 100644
--- a/mitmproxy/certs.py
+++ b/mitmproxy/certs.py
@@ -80,7 +80,7 @@ def create_ca(o, cn, exp):
     return key, cert
 
 
-def dummy_cert(privkey, cacert, commonname, sans):
+def dummy_cert(privkey, cacert, commonname, sans, o):
     """
         Generates a dummy certificate.
 
@@ -88,6 +88,7 @@ def dummy_cert(privkey, cacert, commonname, sans):
         cacert: CA certificate
         commonname: Common name for the generated certificate.
         sans: A list of Subject Alternate Names.
+        o: Organization name for the generated certificate.
 
         Returns cert if operation succeeded, None if not.
     """
@@ -107,6 +108,8 @@ def dummy_cert(privkey, cacert, commonname, sans):
     cert.set_issuer(cacert.get_subject())
     if commonname is not None and len(commonname) < 64:
         cert.get_subject().CN = commonname
+    if o is not None:
+        cert.get_subject().O = o
     cert.set_serial_number(int(time.time() * 10000))
     if ss:
         cert.set_version(2)
@@ -305,7 +308,7 @@ class CertStore:
             ret.append(b"*." + b".".join(parts[i:]))
         return ret
 
-    def get_cert(self, commonname: typing.Optional[bytes], sans: typing.List[bytes]):
+    def get_cert(self, commonname: typing.Optional[bytes], sans: typing.List[bytes], o: typing.Optional[bytes] = None):
         """
             Returns an (cert, privkey, cert_chain) tuple.
 
@@ -313,6 +316,8 @@ class CertStore:
             valid, plain-ASCII, IDNA-encoded domain name.
 
             sans: A list of Subject Alternate Names.
+
+            o: Organization name for the generated certificate.
         """
 
         potential_keys: typing.List[TCertId] = []
@@ -335,7 +340,8 @@ class CertStore:
                     self.default_privatekey,
                     self.default_ca,
                     commonname,
-                    sans),
+                    sans,
+                    o),
                 privatekey=self.default_privatekey,
                 chain_file=self.default_chain_file)
             self.certs[(commonname, tuple(sans))] = entry
@@ -447,6 +453,14 @@ class Cert(serializable.Serializable):
                 c = i[1]
         return c
 
+    @property
+    def o(self):
+        c = None
+        for i in self.subject:
+            if i[0] == b"O":
+                c = i[1]
+        return c
+
     @property
     def altnames(self):
         """
diff --git a/mitmproxy/proxy/protocol/tls.py b/mitmproxy/proxy/protocol/tls.py
index 3f337a2ab..3577af13d 100644
--- a/mitmproxy/proxy/protocol/tls.py
+++ b/mitmproxy/proxy/protocol/tls.py
@@ -469,6 +469,7 @@ class TlsLayer(base.Layer):
         """
         host = None
         sans = set()
+        o = None
 
         # In normal operation, the server address should always be known at this point.
         # However, we may just want to establish TLS so that we can send an error message to the client,
@@ -488,6 +489,8 @@ class TlsLayer(base.Layer):
             if upstream_cert.cn:
                 sans.add(host)
                 host = upstream_cert.cn.decode("utf8").encode("idna")
+            if upstream_cert.o:
+                o = upstream_cert.o
         # Also add SNI values.
         if self._client_hello.sni:
             sans.add(self._client_hello.sni.encode("idna"))
@@ -498,4 +501,4 @@ class TlsLayer(base.Layer):
         # In other words, the Common Name is irrelevant then.
         if host:
             sans.add(host)
-        return self.config.certstore.get_cert(host, list(sans))
+        return self.config.certstore.get_cert(host, list(sans), o)
diff --git a/test/mitmproxy/test_certs.py b/test/mitmproxy/test_certs.py
index 8421ec582..9b4c95163 100644
--- a/test/mitmproxy/test_certs.py
+++ b/test/mitmproxy/test_certs.py
@@ -129,18 +129,22 @@ class TestDummyCert:
             ca.default_privatekey,
             ca.default_ca,
             b"foo.com",
-            [b"one.com", b"two.com", b"*.three.com", b"127.0.0.1"]
+            [b"one.com", b"two.com", b"*.three.com", b"127.0.0.1"],
+            b"Foo Ltd."
         )
         assert r.cn == b"foo.com"
         assert r.altnames == [b'one.com', b'two.com', b'*.three.com']
+        assert r.o == b"Foo Ltd."
 
         r = certs.dummy_cert(
             ca.default_privatekey,
             ca.default_ca,
             None,
-            []
+            [],
+            None
         )
         assert r.cn is None
+        assert r.o is None
         assert r.altnames == []
 
 
@@ -152,6 +156,7 @@ class TestCert:
         c1 = certs.Cert.from_pem(d)
         assert c1.cn == b"google.com"
         assert len(c1.altnames) == 436
+        assert c1.o == b"Google Inc"
 
         with open(tdata.path("mitmproxy/net/data/text_cert_2"), "rb") as f:
             d = f.read()