Add more access to certificate extension data
diff --git a/ChangeLog b/ChangeLog
index 8445143..7472b3e 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,12 @@
+2011-04-06  Jean-Paul Calderone  <exarkun@twistedmatrix.com>
+
+	* OpenSSL/crypto/x509.c: Add get_extension_count and get_extension
+	  to the X509 type, allowing read access to certificate extensions.
+
+	* OpenSSL/crypto/x509ext.c: Add get_short_name and get_data to the
+	  X509Extension type, allowing read access to the contents of an
+	  extension.
+
 2011-03-21  Olivier Hervieu <lp:~ohe>
 
 	* OpenSSL/ssl/ssl.c: Expose a number of symbolic constants for
diff --git a/OpenSSL/crypto/x509.c b/OpenSSL/crypto/x509.c
index 08dc567..5d03d2e 100644
--- a/OpenSSL/crypto/x509.c
+++ b/OpenSSL/crypto/x509.c
@@ -13,6 +13,7 @@
 #include <Python.h>
 #define crypto_MODULE
 #include "crypto.h"
+#include "x509ext.h"
 
 /*
  * X.509 is a standard for digital certificates.  See e.g. the OpenSSL homepage
@@ -300,7 +301,7 @@
 
     py_pkey = crypto_PKey_New(pkey, 1);
     if (py_pkey != NULL) {
-	py_pkey->only_public = 1;
+        py_pkey->only_public = 1;
     }
     return (PyObject *)py_pkey;
 }
@@ -685,6 +686,52 @@
     return Py_None;
 }
 
+static char crypto_X509_get_extension_count_doc[] = "\n\
+Get the number of extensions on the certificate.\n\
+\n\
+@return: Number of extensions as a Python integer\n\
+";
+
+static PyObject *
+crypto_X509_get_extension_count(crypto_X509Obj *self, PyObject *args) {
+    if (!PyArg_ParseTuple(args, ":get_extension_count")) {
+        return NULL;
+    }
+
+    return PyLong_FromLong((long)X509_get_ext_count(self->x509));
+}
+
+static char crypto_X509_get_extension_doc[] = "\n\
+Get a specific extension of the certificate by index.\n\
+\n\
+@param index: The index of the extension to retrieve.\n\
+@return: The X509Extension object at the specified index.\n\
+";
+
+static PyObject *
+crypto_X509_get_extension(crypto_X509Obj *self, PyObject *args) {
+    crypto_X509ExtensionObj *extobj;
+    int loc;
+    X509_EXTENSION *ext;
+
+    if (!PyArg_ParseTuple(args, "i:get_extension", &loc)) {
+        return NULL;
+    }
+
+    /* will return NULL if loc is outside the range of extensions,
+       not registered as an error*/
+    ext = X509_get_ext(self->x509, loc);
+    if (!ext) {
+        PyErr_SetString(PyExc_IndexError, "extension index out of bounds");
+        return NULL; /* Should be reported as an IndexError ? */
+    }
+
+    extobj = PyObject_New(crypto_X509ExtensionObj, &crypto_X509Extension_Type);
+    extobj->x509_extension = X509_EXTENSION_dup(ext);
+
+    return (PyObject*)extobj;
+}
+
 /*
  * ADD_METHOD(name) expands to a correct PyMethodDef declaration
  *   {  'name', (PyCFunction)crypto_X509_name, METH_VARARGS }
@@ -715,6 +762,8 @@
     ADD_METHOD(subject_name_hash),
     ADD_METHOD(digest),
     ADD_METHOD(add_extensions),
+    ADD_METHOD(get_extension),
+    ADD_METHOD(get_extension_count),
     { NULL, NULL }
 };
 #undef ADD_METHOD
diff --git a/OpenSSL/crypto/x509ext.c b/OpenSSL/crypto/x509ext.c
index a874204..bdaac7d 100644
--- a/OpenSSL/crypto/x509ext.c
+++ b/OpenSSL/crypto/x509ext.c
@@ -51,6 +51,26 @@
 }
 
 
+static char crypto_X509Extension_get_data_doc[] = "\n\
+Returns the data of the X509Extension\n\
+\n\
+@return: A C{str} giving the X509Extension's ASN.1 encoded data.\n\
+";
+
+static PyObject *
+crypto_X509Extension_get_data(crypto_X509ExtensionObj *self, PyObject *args) {
+    ASN1_OCTET_STRING *data;
+    PyObject *result;
+
+    if (!PyArg_ParseTuple(args, ":get_data")) {
+        return NULL;
+    }
+
+    data = X509_EXTENSION_get_data(self->x509_extension);
+    result = PyBytes_FromStringAndSize((const char*)data->data, data->length);
+    return result;
+}
+
 /*
  * ADD_METHOD(name) expands to a correct PyMethodDef declaration
  *   {  'name', (PyCFunction)crypto_X509Extension_name, METH_VARARGS }
@@ -62,6 +82,7 @@
 {
     ADD_METHOD(get_critical),
     ADD_METHOD(get_short_name),
+    ADD_METHOD(get_data),
     { NULL, NULL }
 };
 #undef ADD_METHOD
diff --git a/OpenSSL/test/test_crypto.py b/OpenSSL/test/test_crypto.py
index f2831c7..661ee53 100644
--- a/OpenSSL/test/test_crypto.py
+++ b/OpenSSL/test/test_crypto.py
@@ -353,6 +353,26 @@
         self.assertEqual(ext.get_short_name(), b('nsComment'))
 
 
+    def test_get_data(self):
+        """
+        L{X509Extension.get_data} returns a string giving the data of the
+        extension.
+        """
+        ext = X509Extension(b('basicConstraints'), True, b('CA:true'))
+        # Expect to get back the DER encoded form of CA:true.
+        self.assertEqual(ext.get_data(), b('0\x03\x01\x01\xff'))
+
+
+    def test_get_data_wrong_args(self):
+        """
+        L{X509Extension.get_data} raises L{TypeError} if passed any arguments.
+        """
+        ext = X509Extension(b('basicConstraints'), True, b('CA:true'))
+        self.assertRaises(TypeError, ext.get_data, None)
+        self.assertRaises(TypeError, ext.get_data, "foo")
+        self.assertRaises(TypeError, ext.get_data, 7)
+
+
     def test_unused_subject(self):
         """
         The C{subject} parameter to L{X509Extension} may be provided for an
@@ -947,6 +967,26 @@
     """
     pemData = cleartextCertificatePEM + cleartextPrivateKeyPEM
 
+    extpem = """
+-----BEGIN CERTIFICATE-----
+MIIC3jCCAkegAwIBAgIJAJHFjlcCgnQzMA0GCSqGSIb3DQEBBQUAMEcxCzAJBgNV
+BAYTAlNFMRUwEwYDVQQIEwxXZXN0ZXJib3R0b20xEjAQBgNVBAoTCUNhdGFsb2dp
+eDENMAsGA1UEAxMEUm9vdDAeFw0wODA0MjIxNDQ1MzhaFw0wOTA0MjIxNDQ1Mzha
+MFQxCzAJBgNVBAYTAlNFMQswCQYDVQQIEwJXQjEUMBIGA1UEChMLT3Blbk1ldGFk
+aXIxIjAgBgNVBAMTGW5vZGUxLm9tMi5vcGVubWV0YWRpci5vcmcwgZ8wDQYJKoZI
+hvcNAQEBBQADgY0AMIGJAoGBAPIcQMrwbk2nESF/0JKibj9i1x95XYAOwP+LarwT
+Op4EQbdlI9SY+uqYqlERhF19w7CS+S6oyqx0DRZSk4Y9dZ9j9/xgm2u/f136YS1u
+zgYFPvfUs6PqYLPSM8Bw+SjJ+7+2+TN+Tkiof9WP1cMjodQwOmdsiRbR0/J7+b1B
+hec1AgMBAAGjgcQwgcEwCQYDVR0TBAIwADAsBglghkgBhvhCAQ0EHxYdT3BlblNT
+TCBHZW5lcmF0ZWQgQ2VydGlmaWNhdGUwHQYDVR0OBBYEFIdHsBcMVVMbAO7j6NCj
+03HgLnHaMB8GA1UdIwQYMBaAFL2h9Bf9Mre4vTdOiHTGAt7BRY/8MEYGA1UdEQQ/
+MD2CDSouZXhhbXBsZS5vcmeCESoub20yLmV4bWFwbGUuY29thwSC7wgKgRNvbTJA
+b3Blbm1ldGFkaXIub3JnMA0GCSqGSIb3DQEBBQUAA4GBALd7WdXkp2KvZ7/PuWZA
+MPlIxyjS+Ly11+BNE0xGQRp9Wz+2lABtpgNqssvU156+HkKd02rGheb2tj7MX9hG
+uZzbwDAZzJPjzDQDD7d3cWsrVcfIdqVU7epHqIadnOF+X0ghJ39pAm6VVadnSXCt
+WpOdIpB8KksUTCzV591Nr1wd
+-----END CERTIFICATE-----
+    """
     def signable(self):
         """
         Create and return a new L{X509}.
@@ -1199,6 +1239,77 @@
             b("A8:EB:07:F8:53:25:0A:F2:56:05:C5:A5:C4:C4:C7:15"))
 
 
+    def _extcert(self, pkey, extensions):
+        cert = X509()
+        cert.set_pubkey(pkey)
+        cert.get_subject().commonName = "Unit Tests"
+        cert.get_issuer().commonName = "Unit Tests"
+        when = b(datetime.now().strftime("%Y%m%d%H%M%SZ"))
+        cert.set_notBefore(when)
+        cert.set_notAfter(when)
+
+        cert.add_extensions(extensions)
+        return load_certificate(
+            FILETYPE_PEM, dump_certificate(FILETYPE_PEM, cert))
+
+
+    def test_extension_count(self):
+        """
+        L{X509.get_extension_count} returns the number of extensions that are
+        present in the certificate.
+        """
+        pkey = load_privatekey(FILETYPE_PEM, client_key_pem)
+        ca = X509Extension(b('basicConstraints'), True, b('CA:FALSE'))
+        key = X509Extension(b('keyUsage'), True, b('digitalSignature'))
+        subjectAltName = X509Extension(
+            b('subjectAltName'), True, b('DNS:example.com'))
+
+        # Try a certificate with no extensions at all.
+        c = self._extcert(pkey, [])
+        self.assertEqual(c.get_extension_count(), 0)
+
+        # And a certificate with one
+        c = self._extcert(pkey, [ca])
+        self.assertEqual(c.get_extension_count(), 1)
+
+        # And a certificate with several
+        c = self._extcert(pkey, [ca, key, subjectAltName])
+        self.assertEqual(c.get_extension_count(), 3)
+
+
+    def test_get_extension(self):
+        """
+        L{X509.get_extension} takes an integer and returns an L{X509Extension}
+        corresponding to the extension at that index.
+        """
+        pkey = load_privatekey(FILETYPE_PEM, client_key_pem)
+        ca = X509Extension(b('basicConstraints'), True, b('CA:FALSE'))
+        key = X509Extension(b('keyUsage'), True, b('digitalSignature'))
+        subjectAltName = X509Extension(
+            b('subjectAltName'), False, b('DNS:example.com'))
+
+        cert = self._extcert(pkey, [ca, key, subjectAltName])
+
+        ext = cert.get_extension(0)
+        self.assertTrue(isinstance(ext, X509Extension))
+        self.assertTrue(ext.get_critical())
+        self.assertEqual(ext.get_short_name(), b('basicConstraints'))
+
+        ext = cert.get_extension(1)
+        self.assertTrue(isinstance(ext, X509Extension))
+        self.assertTrue(ext.get_critical())
+        self.assertEqual(ext.get_short_name(), b('keyUsage'))
+
+        ext = cert.get_extension(2)
+        self.assertTrue(isinstance(ext, X509Extension))
+        self.assertFalse(ext.get_critical())
+        self.assertEqual(ext.get_short_name(), b('subjectAltName'))
+
+        self.assertRaises(IndexError, cert.get_extension, -1)
+        self.assertRaises(IndexError, cert.get_extension, 4)
+        self.assertRaises(TypeError, cert.get_extension, "hello")
+
+
     def test_invalid_digest_algorithm(self):
         """
         L{X509.digest} raises L{ValueError} if called with an unrecognized hash
diff --git a/OpenSSL/test/util.py b/OpenSSL/test/util.py
index e5b5dc3..f6e9291 100644
--- a/OpenSSL/test/util.py
+++ b/OpenSSL/test/util.py
@@ -15,16 +15,14 @@
 
 from OpenSSL.crypto import Error, _exception_from_error_queue
 
-
-try:
-    bytes = bytes
-except NameError:
+if sys.version_info < (3, 0):
     def b(s):
         return s
     bytes = str
 else:
     def b(s):
-        return s.encode("ascii")
+        return s.encode("charmap")
+    bytes = bytes
 
 
 class TestCase(TestCase):
diff --git a/doc/pyOpenSSL.tex b/doc/pyOpenSSL.tex
index 294008c..7c8bfca 100644
--- a/doc/pyOpenSSL.tex
+++ b/doc/pyOpenSSL.tex
@@ -269,7 +269,7 @@
 
 \begin{funcdesc}{load_crl}{type, buffer}
 Load Certificate Revocation List (CRL) data from a string \var{buffer}.
-\var{buffer} encoded with the type \var{type}.  The type \var{type} 
+\var{buffer} encoded with the type \var{type}.  The type \var{type}
 must either \constant{FILETYPE_PEM} or \constant{FILETYPE_ASN1}).
 \end{funcdesc}
 
@@ -305,6 +305,24 @@
 \versionadded{0.11}
 \end{funcdesc}
 
+\subsubsection{X509Extension objects \label{openssl-x509ext}}
+
+X509Extension objects have the following methods:
+
+\begin{methoddesc}[X509Extension]{get_short_name}{}
+Retrieve the short descriptive name for this extension.
+
+The result is a byte string like \code{``basicConstraints''}.
+\versionadded{0.12}
+\end{methoddesc}
+
+\begin{methoddesc}[X509Extension]{get_data}{}
+Retrieve the data for this extension.
+
+The result is the ASN.1 encoded form of the extension data as a byte string.
+\versionadded{0.12}
+\end{methoddesc}
+
 \subsubsection{X509 objects \label{openssl-x509}}
 
 X509 objects have the following methods:
@@ -424,6 +442,20 @@
 Add the extensions in the sequence \var{extensions} to the certificate.
 \end{methoddesc}
 
+\begin{methoddesc}[X509]{get_extension_count}{}
+Return the number of extensions on this certificate.
+\versionadded{0.12}
+\end{methoddesc}
+
+\begin{methoddesc}[X509]{get_extension}{index}
+Retrieve the extension on this certificate at the given index.
+
+Extensions on a certificate are kept in order.  The index parameter selects
+which extension will be returned.  The returned object will be an X509Extension
+instance.
+\versionadded{0.12}
+\end{methoddesc}
+
 \subsubsection{X509Name objects \label{openssl-x509name}}
 
 X509Name objects have the following methods: