Merge pull request #2267 from queenp/iss2255

added get_extension_for_class #2255

3 files changed, 51 insertions, 0 deletions

diff --git a/docs/x509/reference.rst b/docs/x509/reference.rst
index 8d5d6a6f..62bdb3a9 100644
--- a/docs/x509/reference.rst
+++ b/docs/x509/reference.rst
@@ -889,6 +889,23 @@ X.509 Extensions
             >>> cert.extensions.get_extension_for_oid(ExtensionOID.BASIC_CONSTRAINTS)
             <Extension(oid=<ObjectIdentifier(oid=2.5.29.19, name=basicConstraints)>, critical=True, value=<BasicConstraints(ca=True, path_length=None)>)>
 
+    .. method:: get_extension_for_class(extclass)
+
+        .. versionadded:: 1.1
+
+        :param extclass: An extension class.
+
+        :returns: An instance of the extension class.
+
+        :raises cryptography.x509.ExtensionNotFound: If the certificate does
+            not have the extension requested.
+
+        .. doctest::
+
+            >>> from cryptography import x509
+            >>> cert.extensions.get_extension_for_class(x509.BasicConstraints)
+            <Extension(oid=<ObjectIdentifier(oid=2.5.29.19, name=basicConstraints)>, critical=True, value=<BasicConstraints(ca=True, path_length=None)>)>
+
 .. class:: Extension
 
     .. versionadded:: 0.9
diff --git a/src/cryptography/x509/extensions.py b/src/cryptography/x509/extensions.py
index 798a0e3a..803d7ec5 100644
--- a/src/cryptography/x509/extensions.py
+++ b/src/cryptography/x509/extensions.py
@@ -89,6 +89,15 @@ class Extensions(object):
 
         raise ExtensionNotFound("No {0} extension was found".format(oid), oid)
 
+    def get_extension_for_class(self, extclass):
+        for ext in self:
+            if isinstance(ext.value, extclass):
+                return ext
+
+        raise ExtensionNotFound(
+            "No {0} extension was found".format(extclass), extclass.oid
+        )
+
     def __iter__(self):
         return iter(self._extensions)
 
diff --git a/tests/test_x509_ext.py b/tests/test_x509_ext.py
index 2c5438a9..85373973 100644
--- a/tests/test_x509_ext.py
+++ b/tests/test_x509_ext.py
@@ -832,6 +832,31 @@ class TestExtensions(object):
         extensions = cert.extensions
         assert len(extensions) == 0
 
+    def test_no_extensions_get_for_class(self, backend):
+        cert = _load_cert(
+            os.path.join(
+                "x509", "cryptography.io.pem"
+            ),
+            x509.load_pem_x509_certificate,
+            backend
+        )
+        exts = cert.extensions
+        with pytest.raises(x509.ExtensionNotFound) as exc:
+            exts.get_extension_for_class(x509.IssuerAlternativeName)
+        assert exc.value.oid == ExtensionOID.ISSUER_ALTERNATIVE_NAME
+
+    def test_one_extension_get_for_class(self, backend):
+        cert = _load_cert(
+            os.path.join(
+                "x509", "custom", "basic_constraints_not_critical.pem"
+            ),
+            x509.load_pem_x509_certificate,
+            backend
+        )
+        ext = cert.extensions.get_extension_for_class(x509.BasicConstraints)
+        assert ext is not None
+        assert isinstance(ext.value, x509.BasicConstraints)
+
 
 @pytest.mark.requires_backend_interface(interface=RSABackend)
 @pytest.mark.requires_backend_interface(interface=X509Backend)