Simpler PEM marker code

sybrenstuvel · sybrenstuvel · commit 6ab5251bd62a · 2011-07-21T22:04:39.000+02:00
diff --git a/rsa/key.py b/rsa/key.py
@@ -12,9 +12,6 @@
 import rsa.prime
 import rsa.pem
 
-PEM_PRIVATE_KEY_START = '-----BEGIN RSA PRIVATE KEY-----'
-PEM_PRIVATE_KEY_END = '-----END RSA PRIVATE KEY-----'
-
 
 class PublicKey(object):
     '''Represents a public RSA key.
@@ -339,7 +336,7 @@ def load_private_key_pem(keyfile):
     @return: a PrivateKey object
     '''
 
-    der = rsa.pem.load_pem(keyfile, PEM_PRIVATE_KEY_START, PEM_PRIVATE_KEY_END)
+    der = rsa.pem.load_pem(keyfile, 'RSA PRIVATE KEY')
     return load_private_key_der(der)
 
 def save_private_key_pem(priv_key):
@@ -350,7 +347,7 @@ def save_private_key_pem(priv_key):
     '''
 
     der = save_private_key_der(priv_key)
-    return rsa.pem.save_pem(der, PEM_PRIVATE_KEY_START, PEM_PRIVATE_KEY_END)
+    return rsa.pem.save_pem(der, 'RSA PRIVATE KEY')
 
 
 __all__ = ['PublicKey', 'PrivateKey', 'newkeys', 'load']
diff --git a/rsa/pem.py b/rsa/pem.py
@@ -2,18 +2,24 @@
 
 import base64
 
-def load_pem(contents, pem_start, pem_end):
-    '''Loads a PEM file.
+def _markers(pem_marker):
+    '''Returns the start and end PEM markers
+
+    >>> _markers('RSA PRIVATE KEY')
+    ('-----BEGIN RSA PRIVATE KEY-----', '-----END RSA PRIVATE KEY-----')
 
-    Only considers the information between lines "pem_start" and "pem_end". For
-    private keys these are  '-----BEGIN RSA PRIVATE KEY-----' and 
-    '-----END RSA PRIVATE KEY-----'
+    '''
+
+    return ('-----BEGIN %s-----' % pem_marker,
+            '-----END %s-----' % pem_marker)
+
+def load_pem(contents, pem_marker):
+    '''Loads a PEM file.
 
     @param contents: the contents of the file to interpret
-    @param pem_start: the start marker of the PEM content, such as
-        '-----BEGIN RSA PRIVATE KEY-----'
-    @param pem_end: the end marker of the PEM content, such as
-        '-----END RSA PRIVATE KEY-----'
+    @param pem_marker: the marker of the PEM content, such as 'RSA PRIVATE KEY'
+        when your file has '-----BEGIN RSA PRIVATE KEY-----' and
+        '-----END RSA PRIVATE KEY-----' markers.
 
     @return the base64-decoded content between the start and end markers.
 
@@ -22,6 +28,8 @@ def load_pem(contents, pem_start, pem_end):
 
     '''
 
+    (pem_start, pem_end) = _markers(pem_marker)
+
     pem_lines = []
     in_pem_part = False
 
@@ -58,22 +66,20 @@ def load_pem(contents, pem_start, pem_end):
     pem = ''.join(pem_lines)
     return base64.decodestring(pem)
 
-def save_pem(contents, pem_start, pem_end):
+def save_pem(contents, pem_marker):
     '''Saves a PEM file.
 
-    The PEM file will start with the 'pem_start' marker, then the
-    base64-encoded content, and end with the 'pem_end' marker.
-
     @param contents: the contents to encode in PEM format
-    @param pem_start: the start marker of the PEM content, such as
-        '-----BEGIN RSA PRIVATE KEY-----'
-    @param pem_end: the end marker of the PEM content, such as
-        '-----END RSA PRIVATE KEY-----'
+    @param pem_marker: the marker of the PEM content, such as 'RSA PRIVATE KEY'
+        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.
 
     '''
 
+    (pem_start, pem_end) = _markers(pem_marker)
+
     b64 = base64.encodestring(contents).strip()
     pem_lines = [pem_start]
     
