Ensuring that PEM output is always in bytes.

sybrenstuvel · sybrenstuvel · commit fedcaa1f94c6 · 2016-03-17T15:16:23.000+01:00
This may break some applications. However, it does make the RSA library
easier to use on different Python versions.
diff --git a/rsa/cli.py b/rsa/cli.py
@@ -83,6 +83,8 @@ def keygen():
             outfile.write(data)
     else:
         print('Writing private key to stdout', file=sys.stderr)
+        if sys.version_info[0] >= 3:
+            data = data.decode('ascii')  # on Py3 we must write text, not bytes
         sys.stdout.write(data)
 
 
@@ -189,6 +191,8 @@ def write_outfile(self, outdata, outname):
                 outfile.write(outdata)
         else:
             print('Writing output to stdout', file=sys.stderr)
+            if sys.version_info[0] >= 3:
+                data = outdata.decode('ascii')  # on Py3 we must write text, not bytes
             sys.stdout.write(outdata)
 
 
diff --git a/rsa/pem.py b/rsa/pem.py
@@ -22,14 +22,14 @@
 
 def _markers(pem_marker):
     """
-    Returns the start and end PEM markers
+    Returns the start and end PEM markers, as bytes.
     """
 
-    if is_bytes(pem_marker):
-        pem_marker = pem_marker.decode('utf-8')
+    if not is_bytes(pem_marker):
+        pem_marker = pem_marker.encode('ascii')
 
-    return (b('-----BEGIN %s-----' % pem_marker),
-            b('-----END %s-----' % pem_marker))
+    return (b('-----BEGIN ') + pem_marker + b('-----'),
+            b('-----END ') + pem_marker + b('-----'))
 
 
 def load_pem(contents, pem_marker):
@@ -106,7 +106,7 @@ def save_pem(contents, pem_marker):
         when your file has '-----BEGIN RSA PRIVATE KEY-----' and
         '-----END RSA PRIVATE KEY-----' markers.
 
-    :return: the base64-encoded content between the start and end markers.
+    :return: the base64-encoded content between the start and end markers, as bytes.
 
     """
 
diff --git a/tests/test_pem.py b/tests/test_pem.py
@@ -17,7 +17,7 @@
 
 import unittest
 
-from rsa._compat import b
+from rsa._compat import b, is_bytes
 from rsa.pem import _markers
 import rsa.key
 
@@ -72,3 +72,17 @@ def test_bytes_private(self):
         key = rsa.key.PrivateKey.load_pkcs1(private_key_pem.encode('ascii'))
         self.assertEqual(prime1, key.p)
         self.assertEqual(prime2, key.q)
+
+
+class TestByteOutput(unittest.TestCase):
+    """Tests that PEM and DER are returned as bytes."""
+
+    def test_bytes_public(self):
+        key = rsa.key.PublicKey.load_pkcs1_openssl_pem(public_key_pem)
+        self.assertTrue(is_bytes(key.save_pkcs1(format='DER')))
+        self.assertTrue(is_bytes(key.save_pkcs1(format='PEM')))
+
+    def test_bytes_private(self):
+        key = rsa.key.PrivateKey.load_pkcs1(private_key_pem)
+        self.assertTrue(is_bytes(key.save_pkcs1(format='DER')))
+        self.assertTrue(is_bytes(key.save_pkcs1(format='PEM')))
