Added loading of DER and PEM encoded private keys

3 files changed, 188 insertions, 14 deletions

diff --git a/rsa/key.py b/rsa/key.py
index fd4f4db..0fe4a89 100644
--- a/rsa/key.py
+++ b/rsa/key.py
@@ -6,6 +6,7 @@ PrivateKey object.
 '''
 
 import rsa.prime
+import rsa.pem
 
 class PublicKey(object):
     '''Represents a public RSA key.
@@ -75,19 +76,6 @@ class PrivateKey(object):
 
     '''
 
-    # RSAPrivateKey ::= SEQUENCE {
-    #     version           Version, 
-    #     modulus           INTEGER,  -- n
-    #     publicExponent    INTEGER,  -- e
-    #     privateExponent   INTEGER,  -- d
-    #     prime1            INTEGER,  -- p
-    #     prime2            INTEGER,  -- q
-    #     exponent1         INTEGER,  -- d mod (p-1)
-    #     exponent2         INTEGER,  -- d mod (q-1) 
-    #     coefficient       INTEGER,  -- (inverse of q) mod p
-    #     otherPrimeInfos   OtherPrimeInfos OPTIONAL 
-    # }
-
     __slots__ = ('n', 'e', 'd', 'p', 'q', 'exp1', 'exp2', 'coef')
 
     def __init__(self, n, e, d, p, q, exp1=None, exp2=None, coef=None):
@@ -119,6 +107,24 @@ class PrivateKey(object):
     def __repr__(self):
         return u'PrivateKey(%(n)i, %(e)i, %(d)i, %(p)i, %(q)i)' % self
 
+    def __eq__(self, other):
+        if other is None:
+            return False
+
+        if not isinstance(other, PrivateKey):
+            return False
+
+        return (self.n == other.n and
+            self.e == other.e and
+            self.d == other.d and
+            self.p == other.p and
+            self.q == other.q and
+            self.exp1 == other.exp1 and
+            self.exp2 == other.exp2 and
+            self.coef == other.coef)
+
+    def __ne__(self, other):
+        return not (self == other)
 
 def extended_gcd(a, b):
     """Returns a tuple (r, i, j) such that r = gcd(a, b) = ia + jb
@@ -234,7 +240,68 @@ def newkeys(nbits):
         PrivateKey(n, e, d, p, q)
     )
 
-__all__ = ['PublicKey', 'PrivateKey', 'newkeys']
+def load_private_key_der(keyfile):
+    r'''Loads a key in DER format.
+
+    @param keyfile: contents of a DER-encoded file that contains the private
+        key.
+    @return: a PrivateKey object
+
+    First let's construct a DER encoded key:
+
+    >>> import base64
+    >>> b64der = 'MC4CAQACBQDeKYlRAgMBAAECBQDHn4npAgMA/icCAwDfxwIDANcXAgInbwIDAMZt'
+    >>> der = base64.decodestring(b64der)
+
+    This loads the file:
+
+    >>> load_private_key_der(der)
+    PrivateKey(3727264081, 65537, 3349121513, 65063, 57287)
+
+    '''
+
+    from pyasn1.codec.der import decoder
+    (priv, _) = decoder.decode(keyfile)
+
+    # ASN.1 contents of DER encoded private key:
+    #
+    # RSAPrivateKey ::= SEQUENCE {
+    #     version           Version, 
+    #     modulus           INTEGER,  -- n
+    #     publicExponent    INTEGER,  -- e
+    #     privateExponent   INTEGER,  -- d
+    #     prime1            INTEGER,  -- p
+    #     prime2            INTEGER,  -- q
+    #     exponent1         INTEGER,  -- d mod (p-1)
+    #     exponent2         INTEGER,  -- d mod (q-1) 
+    #     coefficient       INTEGER,  -- (inverse of q) mod p
+    #     otherPrimeInfos   OtherPrimeInfos OPTIONAL 
+    # }
+
+    if priv[0] != 0:
+        raise ValueError('Unable to read this file, version %s != 0' % priv[0])
+
+    return PrivateKey(*priv[1:9])
+
+def load_private_key_pem(keyfile):
+    '''Loads a PEM-encoded private key file.
+
+    The contents of the file before the "-----BEGIN RSA PRIVATE KEY-----" and
+    after the "-----END RSA PRIVATE KEY-----" lines is ignored.
+
+    @param keyfile: contents of a PEM-encoded file that contains the private
+        key.
+    @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)
+    return load_private_key_der(der)
+
+
+__all__ = ['PublicKey', 'PrivateKey', 'newkeys', 'load']
 
 if __name__ == '__main__':
     import doctest
diff --git a/rsa/pem.py b/rsa/pem.py
new file mode 100644
index 0000000..a06d59f
--- /dev/null
+++ b/rsa/pem.py
@@ -0,0 +1,60 @@
+'''Functions that load and write PEM-encoded files.'''
+
+import base64
+
+def load_pem(contents, pem_start, pem_end):
+    '''Loads a PEM file.
+
+    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-----'
+
+    @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-----'
+
+    @return the base64-decoded content between the start and end markers.
+
+    @raise ValueError: when the content is invalid, for example when the start
+        marker cannot be found.
+
+    '''
+
+    pem_lines = []
+    in_pem_part = False
+
+    for line in contents.split('\n'):
+        line = line.strip()
+
+        # Handle start marker
+        if line == pem_start:
+            if in_pem_part:
+                raise ValueError('Seen start marker "%s" twice' % pem_start)
+
+            in_pem_part = True
+            continue
+
+        # Skip stuff before first marker
+        if not in_pem_part:
+            continue
+
+        # Handle end marker
+        if in_pem_part and line == pem_end:
+            in_pem_part = False
+            break
+
+        pem_lines.append(line)
+
+    # Do some sanity checks
+    if not pem_lines:
+        raise ValueError('No PEM start marker "%s" found' % pem_start)
+
+    if in_pem_part:
+        raise ValueError('No PEM end marker "%s" found' % pem_end)
+
+    # Base64-decode the contents
+    pem = ''.join(pem_lines)
+    return base64.decodestring(pem)
+
diff --git a/tests/test_load_save_keys.py b/tests/test_load_save_keys.py
new file mode 100644
index 0000000..695de72
--- /dev/null
+++ b/tests/test_load_save_keys.py
@@ -0,0 +1,47 @@
+'''Unittest for saving and loading keys.'''
+
+import base64
+import unittest
+
+import rsa.key
+
+B64PRIV_DER = 'MC4CAQACBQDeKYlRAgMBAAECBQDHn4npAgMA/icCAwDfxwIDANcXAgInbwIDAMZt'
+PRIVATE_DER = base64.decodestring(B64PRIV_DER)
+
+PRIVATE_PEM = '''
+-----BEGIN CONFUSING STUFF-----
+Cruft before the key
+
+-----BEGIN RSA PRIVATE KEY-----
+%s
+-----END RSA PRIVATE KEY-----
+
+Stuff after the key
+-----END CONFUSING STUFF-----
+''' % B64PRIV_DER
+
+
+class DerTest(unittest.TestCase):
+    '''Test saving and loading DER keys.'''
+
+    def test_load_private_key(self):
+        '''Test loading private DER keys.'''
+
+        key = rsa.key.load_private_key_der(PRIVATE_DER)
+        expected = rsa.key.PrivateKey(3727264081, 65537, 3349121513, 65063, 57287)
+
+        self.assertEqual(expected, key)
+
+
+class PemTest(unittest.TestCase):
+    '''Test saving and loading PEM keys.'''
+
+
+    def test_load_private_key(self):
+        '''Test loading private PEM files.'''
+
+        key = rsa.key.load_private_key_pem(PRIVATE_PEM)
+        expected = rsa.key.PrivateKey(3727264081, 65537, 3349121513, 65063, 57287)
+
+        self.assertEqual(expected, key)
+