From 6ab5251bd62a3330dd495378eeef22e3020fb32c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Thu, 21 Jul 2011 22:04:39 +0200
Subject: Simpler PEM marker code

---
 rsa/pem.py | 40 +++++++++++++++++++++++-----------------
 1 file changed, 23 insertions(+), 17 deletions(-)

(limited to 'rsa/pem.py')

diff --git a/rsa/pem.py b/rsa/pem.py
index 3f0ea76..915a5c4 100644
--- 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]
     
-- 
cgit v1.2.1