simplify the CSRBuilder

2 files changed, 9 insertions, 23 deletions

diff --git a/src/cryptography/hazmat/backends/openssl/backend.py b/src/cryptography/hazmat/backends/openssl/backend.py
index faa3ee55..2712abcb 100644
--- a/src/cryptography/hazmat/backends/openssl/backend.py
+++ b/src/cryptography/hazmat/backends/openssl/backend.py
@@ -1193,15 +1193,10 @@ 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]
+                pp, r = encode(self, extension.value)
+            except KeyError:
                 raise NotImplementedError('Extension not yet supported.')
 
             obj = _txt2obj_gc(self, extension.oid.dotted_string)
diff --git a/src/cryptography/x509.py b/src/cryptography/x509.py
index 7b1de8b8..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: