Merge pull request #2231 from reaperhulk/simplify-certbuilder

Simplify Cert and CSR Builder

4 files changed, 44 insertions, 100 deletions

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 738c0552..f7e04f7e 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -37,15 +37,6 @@ Changelog
 
 * Extension support was added to
   :class:`~cryptography.x509.CertificateSigningRequest`.
-* Add support for creating certificate signing requests with
-  :class:`~cryptography.x509.CertificateSigningRequestBuilder`. This includes
-  support for the following extensions:
-
-  * :class:`~cryptography.x509.BasicConstraints`
-  * :class:`~cryptography.x509.ExtendedKeyUsage`
-  * :class:`~cryptography.x509.KeyUsage`
-  * :class:`~cryptography.x509.SubjectAlternativeName`
-
 * Add support for creating signed certificates with
   :class:`~cryptography.x509.CertificateBuilder`. This includes support for
   the following extensions:
@@ -60,6 +51,10 @@ Changelog
   * :class:`~cryptography.x509.CRLDistributionPoints`
   * :class:`~cryptography.x509.InhibitAnyPolicy`
 
+* Add support for creating certificate signing requests with
+  :class:`~cryptography.x509.CertificateSigningRequestBuilder`. This includes
+  support for the same extensions supported in the ``CertificateBuilder``.
+
 0.9.3 - 2015-07-09
 ~~~~~~~~~~~~~~~~~~
 
diff --git a/src/cryptography/hazmat/backends/openssl/backend.py b/src/cryptography/hazmat/backends/openssl/backend.py
index c583214d..ad88dd9d 100644
--- a/src/cryptography/hazmat/backends/openssl/backend.py
+++ b/src/cryptography/hazmat/backends/openssl/backend.py
@@ -473,6 +473,21 @@ def _encode_crl_distribution_points(backend, crl_distribution_points):
     return pp, r
 
 
+_EXTENSION_ENCODE_HANDLERS = {
+    x509.OID_BASIC_CONSTRAINTS: _encode_basic_constraints,
+    x509.OID_SUBJECT_KEY_IDENTIFIER: _encode_subject_key_identifier,
+    x509.OID_KEY_USAGE: _encode_key_usage,
+    x509.OID_SUBJECT_ALTERNATIVE_NAME: _encode_subject_alt_name,
+    x509.OID_EXTENDED_KEY_USAGE: _encode_extended_key_usage,
+    x509.OID_AUTHORITY_KEY_IDENTIFIER: _encode_authority_key_identifier,
+    x509.OID_AUTHORITY_INFORMATION_ACCESS: (
+        _encode_authority_information_access
+    ),
+    x509.OID_CRL_DISTRIBUTION_POINTS: _encode_crl_distribution_points,
+    x509.OID_INHIBIT_ANY_POLICY: _encode_inhibit_any_policy,
+}
+
+
 @utils.register_interface(CipherBackend)
 @utils.register_interface(CMACBackend)
 @utils.register_interface(DERSerializationBackend)
@@ -1178,17 +1193,12 @@ class Backend(object):
             self._lib.sk_X509_EXTENSION_free,
         )
         for extension in builder._extensions:
-            if isinstance(extension.value, x509.BasicConstraints):
-                pp, r = _encode_basic_constraints(self, extension.value)
-            elif isinstance(extension.value, x509.SubjectAlternativeName):
-                pp, r = _encode_subject_alt_name(self, extension.value)
-            elif isinstance(extension.value, x509.KeyUsage):
-                pp, r = _encode_key_usage(self, extension.value)
-            elif isinstance(extension.value, x509.ExtendedKeyUsage):
-                pp, r = _encode_extended_key_usage(self, extension.value)
-            else:
+            try:
+                encode = _EXTENSION_ENCODE_HANDLERS[extension.oid]
+            except KeyError:
                 raise NotImplementedError('Extension not yet supported.')
 
+            pp, r = encode(self, extension.value)
             obj = _txt2obj_gc(self, extension.oid.dotted_string)
             extension = self._lib.X509_EXTENSION_create_by_OBJ(
                 self._ffi.NULL,
@@ -1279,31 +1289,12 @@ class Backend(object):
 
         # Add extensions.
         for i, extension in enumerate(builder._extensions):
-            if isinstance(extension.value, x509.BasicConstraints):
-                pp, r = _encode_basic_constraints(self, extension.value)
-            elif isinstance(extension.value, x509.AuthorityKeyIdentifier):
-                pp, r = _encode_authority_key_identifier(self, extension.value)
-            elif isinstance(extension.value, x509.KeyUsage):
-                pp, r = _encode_key_usage(self, extension.value)
-            elif isinstance(extension.value, x509.InhibitAnyPolicy):
-                pp, r = _encode_inhibit_any_policy(self, extension.value)
-            elif isinstance(extension.value, x509.ExtendedKeyUsage):
-                pp, r = _encode_extended_key_usage(self, extension.value)
-            elif isinstance(extension.value, x509.SubjectAlternativeName):
-                pp, r = _encode_subject_alt_name(self, extension.value)
-            elif isinstance(extension.value, x509.SubjectKeyIdentifier):
-                pp, r = _encode_subject_key_identifier(self, extension.value)
-            elif isinstance(extension.value, x509.AuthorityInformationAccess):
-                pp, r = _encode_authority_information_access(
-                    self, extension.value
-                )
-            elif isinstance(extension.value, x509.CRLDistributionPoints):
-                pp, r = _encode_crl_distribution_points(
-                    self, extension.value
-                )
-            else:
+            try:
+                encode = _EXTENSION_ENCODE_HANDLERS[extension.oid]
+            except KeyError:
                 raise NotImplementedError('Extension not yet supported.')
 
+            pp, r = encode(self, extension.value)
             obj = _txt2obj_gc(self, extension.oid.dotted_string)
             extension = self._lib.X509_EXTENSION_create_by_OBJ(
                 self._ffi.NULL,
diff --git a/src/cryptography/x509.py b/src/cryptography/x509.py
index bcda7217..a1d0b2f9 100644
--- a/src/cryptography/x509.py
+++ b/src/cryptography/x509.py
@@ -1665,20 +1665,11 @@ class CertificateSigningRequestBuilder(object):
         """
         Adds an X.509 extension to the certificate request.
         """
-        if isinstance(extension, BasicConstraints):
-            extension = Extension(OID_BASIC_CONSTRAINTS, critical, extension)
-        elif isinstance(extension, ExtendedKeyUsage):
-            extension = Extension(OID_EXTENDED_KEY_USAGE, critical, extension)
-        elif isinstance(extension, SubjectAlternativeName):
-            extension = Extension(
-                OID_SUBJECT_ALTERNATIVE_NAME, critical, extension
-            )
-        elif isinstance(extension, KeyUsage):
-            extension = Extension(OID_KEY_USAGE, critical, extension)
-        elif isinstance(extension, InhibitAnyPolicy):
-            extension = Extension(OID_INHIBIT_ANY_POLICY, critical, extension)
-        else:
-            raise NotImplementedError('Unsupported X.509 extension.')
+        if not isinstance(extension, ExtensionType):
+            raise TypeError("extension must be an ExtensionType")
+
+        extension = Extension(extension.oid, critical, extension)
+
         # TODO: This is quadratic in the number of extensions
         for e in self._extensions:
             if e.oid == extension.oid:
@@ -1810,40 +1801,10 @@ class CertificateBuilder(object):
         """
         Adds an X.509 extension to the certificate.
         """
-        if isinstance(extension, BasicConstraints):
-            extension = Extension(OID_BASIC_CONSTRAINTS, critical, extension)
-        elif isinstance(extension, AuthorityKeyIdentifier):
-            extension = Extension(
-                OID_AUTHORITY_KEY_IDENTIFIER, critical, extension
-            )
-        elif isinstance(extension, KeyUsage):
-            extension = Extension(OID_KEY_USAGE, critical, extension)
-        elif isinstance(extension, InhibitAnyPolicy):
-            extension = Extension(OID_INHIBIT_ANY_POLICY, critical, extension)
-        elif isinstance(extension, ExtendedKeyUsage):
-            extension = Extension(OID_EXTENDED_KEY_USAGE, critical, extension)
-        elif isinstance(extension, SubjectAlternativeName):
-            extension = Extension(
-                OID_SUBJECT_ALTERNATIVE_NAME, critical, extension
-            )
-        elif isinstance(extension, AuthorityInformationAccess):
-            extension = Extension(
-                OID_AUTHORITY_INFORMATION_ACCESS, critical, extension
-            )
-        elif isinstance(extension, SubjectKeyIdentifier):
-            extension = Extension(
-                OID_SUBJECT_KEY_IDENTIFIER, critical, extension
-            )
-        elif isinstance(extension, CRLDistributionPoints):
-            extension = Extension(
-                OID_CRL_DISTRIBUTION_POINTS, critical, extension
-            )
-        elif isinstance(extension, IssuerAlternativeName):
-            extension = Extension(
-                OID_ISSUER_ALTERNATIVE_NAME, critical, extension
-            )
-        else:
-            raise NotImplementedError('Unsupported X.509 extension.')
+        if not isinstance(extension, ExtensionType):
+            raise TypeError("extension must be an ExtensionType")
+
+        extension = Extension(extension.oid, critical, extension)
 
         # TODO: This is quadratic in the number of extensions
         for e in self._extensions:
diff --git a/tests/test_x509.py b/tests/test_x509.py
index b630e337..26bd3cb8 100644
--- a/tests/test_x509.py
+++ b/tests/test_x509.py
@@ -1064,10 +1064,10 @@ class TestCertificateBuilder(object):
                 x509.BasicConstraints(ca=False, path_length=None), True,
             )
 
-    def test_add_unsupported_extension(self):
+    def test_add_invalid_extension_type(self):
         builder = x509.CertificateBuilder()
 
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(TypeError):
             builder.add_extension(object(), False)
 
     @pytest.mark.requires_backend_interface(interface=RSABackend)
@@ -1701,15 +1701,13 @@ class TestCertificateSigningRequestBuilder(object):
         with pytest.raises(TypeError):
             builder.subject_name('NotAName')
 
-    def test_add_unsupported_extension(self):
+    def test_add_invalid_extension_type(self):
         builder = x509.CertificateSigningRequestBuilder()
-        with pytest.raises(NotImplementedError):
-            builder.add_extension(
-                x509.AuthorityKeyIdentifier('keyid', None, None),
-                critical=False,
-            )
 
-    def test_add_unsupported_extension_in_backend(self, backend):
+        with pytest.raises(TypeError):
+            builder.add_extension(object(), False)
+
+    def test_add_unsupported_extension(self, backend):
         private_key = RSA_KEY_2048.private_key(backend)
         builder = x509.CertificateSigningRequestBuilder()
         builder = builder.subject_name(
@@ -1720,8 +1718,7 @@ class TestCertificateSigningRequestBuilder(object):
             x509.SubjectAlternativeName([x509.DNSName(u"cryptography.io")]),
             critical=False,
         ).add_extension(
-            x509.InhibitAnyPolicy(0),
-            critical=False
+            x509.IssuerAlternativeName([x509.DNSName(u"crypto.io")]), False
         )
         with pytest.raises(NotImplementedError):
             builder.sign(private_key, hashes.SHA256(), backend)