From c2af63b4e064d9b64fa2b7f65e5c4dda34759d60 Mon Sep 17 00:00:00 2001
From: Danielle Madeley <danielle.madeley@squareweave.com.au>
Date: Mon, 4 Sep 2017 14:17:39 +1000
Subject: [PATCH] Use PEM armor/unarmor functions

---
 pkcs11/util/x509.py |  6 ++++--
 tests/test_x509.py  | 21 ++++++++-------------
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/pkcs11/util/x509.py b/pkcs11/util/x509.py
index c93a662..35d4bd4 100644
--- a/pkcs11/util/x509.py
+++ b/pkcs11/util/x509.py
@@ -2,8 +2,6 @@
 Certificate handling utilities for X.509 (SSL) certificates.
 """
 
-from datetime import datetime
-
 from asn1crypto.x509 import Certificate
 
 from ..constants import Attribute, ObjectClass, CertificateType
@@ -15,6 +13,8 @@ def decode_x509_public_key(der):
     Decode a DER-encoded X.509 certificate's public key into a set of
     attributes able to be passed to :meth:`pkcs11.Session.create_object`.
 
+    For PEM-encoded certificates, use :func:`asn1crypto.pem.unarmor`.
+
     .. warning::
 
         Does not verify certificate.
@@ -69,6 +69,8 @@ def decode_x509_certificate(der, extended_set=False):
     Optionally pass `extended_set` to include additional attributes:
     start date, end date and key identifiers.
 
+    For PEM-encoded certificates, use :func:`asn1crypto.pem.unarmor`.
+
     .. warning::
 
         Does not verify certificate.
diff --git a/tests/test_x509.py b/tests/test_x509.py
index 91c6f87..056e85e 100644
--- a/tests/test_x509.py
+++ b/tests/test_x509.py
@@ -6,7 +6,10 @@
 import subprocess
 import datetime
 
-from asn1crypto.x509 import Certificate
+from asn1crypto import pem
+from asn1crypto.x509 import Certificate, TbsCertificate, Time, Name
+from asn1crypto.keys import RSAPublicKey
+from asn1crypto.csr import CertificationRequest, CertificationRequestInfo
 
 import pkcs11
 from pkcs11.util.rsa import encode_rsa_public_key
@@ -29,7 +32,8 @@
 #   -out cert.pem \
 #   -days 365 \
 #   -nodes
-CERT = base64.b64decode("""
+_, _, CERT = pem.unarmor(b"""
+-----BEGIN CERTIFICATE-----
 MIICKzCCAdWgAwIBAgIJAK3BO9rnLZd9MA0GCSqGSIb3DQEBBQUAMEUxCzAJBgNV
 BAYTAkFVMRMwEQYDVQQIEwpTb21lLVN0YXRlMSEwHwYDVQQKExhJbnRlcm5ldCBX
 aWRnaXRzIFB0eSBMdGQwHhcNMTcwNjAyMDI0ODMyWhcNMTgwNjAyMDI0ODMyWjBF
@@ -42,6 +46,7 @@
 cm5ldCBXaWRnaXRzIFB0eSBMdGSCCQCtwTva5y2XfTAMBgNVHRMEBTADAQH/MA0G
 CSqGSIb3DQEBBQUAA0EAOdvMKLrIFOYF3aVLGharY196heO0fndm39sZAXJ4PItx
 n28DytHEdAoltksfJ2Ds3XAjQqcpI5eBbhIoN9Ckxg==
+-----END CERTIFICATE-----
 """)
 
 
@@ -163,9 +168,6 @@ def test_self_sign_certificate(self):
         # Warning: proof of concept code only!
         pub, priv = self.session.generate_keypair(KeyType.RSA, 1024)
 
-        from asn1crypto.x509 import TbsCertificate, Time, Name
-        from asn1crypto.keys import RSAPublicKey
-
         tbs = TbsCertificate({
             'version': 'v1',
             'serial_number': 1,
@@ -214,9 +216,7 @@ def test_self_sign_certificate(self):
                               stdin=subprocess.PIPE,
                               stdout=subprocess.DEVNULL) as proc:
 
-            proc.stdin.write(b'-----BEGIN CERTIFICATE-----\n')
-            proc.stdin.write(base64.encodebytes(cert.dump()))
-            proc.stdin.write(b'-----END CERTIFICATE-----\n')
+            proc.stdin.write(pem.armor('CERTIFICATE', cert.dump()))
             proc.stdin.close()
 
             self.assertEqual(proc.wait(), 0)
@@ -226,11 +226,6 @@ def test_sign_csr(self):
         # Warning: proof of concept code only!
         pub, priv = self.session.generate_keypair(KeyType.RSA, 1024)
 
-        from asn1crypto.csr import (CertificationRequest,
-                                    CertificationRequestInfo)
-        from asn1crypto.x509 import Name
-        from asn1crypto.keys import RSAPublicKey
-
         info = CertificationRequestInfo({
             'version': 0,
             'subject': Name.build({