Added support for saving private keys in DER and PEM format

4 files changed, 107 insertions, 4 deletions

diff --git a/CHANGELOG.txt b/CHANGELOG.txt
index ce3d3e5..a06b166 100644
--- a/CHANGELOG.txt
+++ b/CHANGELOG.txt
@@ -22,6 +22,8 @@ Version 3.0 - in development
 
 - Modeling private and public key as real objects rather than dicts.
 
+- Support for saving and loading private keys as PEM and DER files
+
 Version 2.0
 ----------------------------------------
 
diff --git a/rsa/key.py b/rsa/key.py
index 0fe4a89..6f348eb 100644
--- a/rsa/key.py
+++ b/rsa/key.py
@@ -3,11 +3,19 @@
 Create new keys with the newkeys() function. It will give you a PublicKey and a
 PrivateKey object.
 
+Loading and saving keys requires the pyasn1 module. This module is imported as
+late as possible, such that other functionality will remain working in absence
+of pyasn1.
+
 '''
 
 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.
 
@@ -283,6 +291,43 @@ def load_private_key_der(keyfile):
 
     return PrivateKey(*priv[1:9])
 
+def save_private_key_der(priv_key):
+    '''Saves the private key in DER format.
+
+    @param priv_key: the private key to save
+    @returns: the DER-encoded private key.
+    '''
+
+    from pyasn1.type import univ, namedtype, tag
+    from pyasn1.codec.der import encoder
+
+    class AsnPrivKey(univ.Sequence):
+        componentType = namedtype.NamedTypes(
+            namedtype.NamedType('version', univ.Integer()),
+            namedtype.NamedType('modulus', univ.Integer()),
+            namedtype.NamedType('publicExponent', univ.Integer()),
+            namedtype.NamedType('privateExponent', univ.Integer()),
+            namedtype.NamedType('prime1', univ.Integer()),
+            namedtype.NamedType('prime2', univ.Integer()),
+            namedtype.NamedType('exponent1', univ.Integer()),
+            namedtype.NamedType('exponent2', univ.Integer()),
+            namedtype.NamedType('coefficient', univ.Integer()),
+        )
+
+    # Create the ASN object
+    asn_key = AsnPrivKey()
+    asn_key.setComponentByName('version', 0)
+    asn_key.setComponentByName('modulus', priv_key.n)
+    asn_key.setComponentByName('publicExponent', priv_key.e)
+    asn_key.setComponentByName('privateExponent', priv_key.d)
+    asn_key.setComponentByName('prime1', priv_key.p)
+    asn_key.setComponentByName('prime2', priv_key.q)
+    asn_key.setComponentByName('exponent1', priv_key.exp1)
+    asn_key.setComponentByName('exponent2', priv_key.exp2)
+    asn_key.setComponentByName('coefficient', priv_key.coef)
+
+    return encoder.encode(asn_key)
+
 def load_private_key_pem(keyfile):
     '''Loads a PEM-encoded private key file.
 
@@ -294,12 +339,19 @@ def load_private_key_pem(keyfile):
     @return: a PrivateKey object
     '''
 
-    PEM_START = '-----BEGIN RSA PRIVATE KEY-----'
-    PEM_END = '-----END RSA PRIVATE KEY-----'
-
-    der = rsa.pem.load_pem(keyfile, PEM_START, PEM_END)
+    der = rsa.pem.load_pem(keyfile, PEM_PRIVATE_KEY_START, PEM_PRIVATE_KEY_END)
     return load_private_key_der(der)
 
+def save_private_key_pem(priv_key):
+    '''Saves a PEM-encoded private key file.
+
+    @param keyfile: a PrivateKey object
+    @return: contents of a PEM-encoded file that contains the private key.
+    '''
+
+    der = save_private_key_der(priv_key)
+    return rsa.pem.save_pem(der, PEM_PRIVATE_KEY_START, PEM_PRIVATE_KEY_END)
+
 
 __all__ = ['PublicKey', 'PrivateKey', 'newkeys', 'load']
 
diff --git a/rsa/pem.py b/rsa/pem.py
index a06d59f..3f0ea76 100644
--- a/rsa/pem.py
+++ b/rsa/pem.py
@@ -58,3 +58,31 @@ def load_pem(contents, pem_start, pem_end):
     pem = ''.join(pem_lines)
     return base64.decodestring(pem)
 
+def save_pem(contents, pem_start, pem_end):
+    '''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-----'
+
+    @return the base64-encoded content between the start and end markers.
+
+    '''
+
+    b64 = base64.encodestring(contents).strip()
+    pem_lines = [pem_start]
+    
+    for block_start in range(0, len(b64), 64):
+        block = b64[block_start:block_start + 64]
+        pem_lines.append(block)
+
+    pem_lines.append(pem_end)
+    pem_lines.append('')
+
+    return '\n'.join(pem_lines)
+    
diff --git a/tests/test_load_save_keys.py b/tests/test_load_save_keys.py
index 695de72..c13d059 100644
--- a/tests/test_load_save_keys.py
+++ b/tests/test_load_save_keys.py
@@ -20,6 +20,12 @@ Stuff after the key
 -----END CONFUSING STUFF-----
 ''' % B64PRIV_DER
 
+CLEAN_PRIVATE_PEM = '''\
+-----BEGIN RSA PRIVATE KEY-----
+%s
+-----END RSA PRIVATE KEY-----
+''' % B64PRIV_DER
+
 
 class DerTest(unittest.TestCase):
     '''Test saving and loading DER keys.'''
@@ -32,6 +38,13 @@ class DerTest(unittest.TestCase):
 
         self.assertEqual(expected, key)
 
+    def test_save_private_key(self):
+        '''Test saving private DER keys.'''
+
+        key = rsa.key.PrivateKey(3727264081, 65537, 3349121513, 65063, 57287)
+        der = rsa.key.save_private_key_der(key)
+
+        self.assertEqual(PRIVATE_DER, der)
 
 class PemTest(unittest.TestCase):
     '''Test saving and loading PEM keys.'''
@@ -45,3 +58,11 @@ class PemTest(unittest.TestCase):
 
         self.assertEqual(expected, key)
 
+    def test_save_private_key(self):
+        '''Test saving private PEM files.'''
+
+        key = rsa.key.PrivateKey(3727264081, 65537, 3349121513, 65063, 57287)
+        pem = rsa.key.save_private_key_pem(key)
+
+        self.assertEqual(CLEAN_PRIVATE_PEM, pem)
+