From ac2d48aefb7b6ec801dcfed8385e9d674b8e2414 Mon Sep 17 00:00:00 2001
From: "Sybren A. St?vel" <sybren@stuvel.eu>
Date: Sun, 24 Jul 2011 18:29:17 +0200
Subject: Added saving and loading public keys in PKCS#1 format (PEM+DER)

---
 rsa/key.py                   | 38 ++++++++++++++++++++++++++++++-
 tests/test_load_save_keys.py | 54 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/rsa/key.py b/rsa/key.py
index 231b22a..2aa39f1 100644
--- a/rsa/key.py
+++ b/rsa/key.py
@@ -68,6 +68,18 @@ class PublicKey(object):
     def __repr__(self):
         return u'PublicKey(%i, %i)' % (self.n, self.e)
 
+    def __eq__(self, other):
+        if other is None:
+            return False
+
+        if not isinstance(other, PublicKey):
+            return False
+
+        return self.n == other.n and self.e == other.e
+
+    def __ne__(self, other):
+        return not (self == other)
+
     @classmethod
     def load_pkcs1_der(cls, keyfile):
         r'''Loads a key in PKCS#1 DER format.
@@ -122,6 +134,31 @@ class PublicKey(object):
 
         return encoder.encode(asn_key)
 
+    @classmethod
+    def load_pkcs1_pem(cls, keyfile):
+        '''Loads a PKCS#1 PEM-encoded public key file.
+
+        The contents of the file before the "-----BEGIN RSA PUBLIC KEY-----" and
+        after the "-----END RSA PUBLIC KEY-----" lines is ignored.
+
+        @param keyfile: contents of a PEM-encoded file that contains the public
+            key.
+        @return: a PublicKey object
+        '''
+
+        der = rsa.pem.load_pem(keyfile, 'RSA PUBLIC KEY')
+        return cls.load_pkcs1_der(der)
+
+    def save_pkcs1_pem(self):
+        '''Saves a PKCS#1 PEM-encoded public key file.
+
+        @return: contents of a PEM-encoded file that contains the public key.
+        '''
+
+        der = self.save_pkcs1_der()
+        return rsa.pem.save_pem(der, 'RSA PUBLIC KEY')
+
+
 class PrivateKey(object):
     '''Represents a private RSA key.
 
@@ -303,7 +340,6 @@ class PrivateKey(object):
     def save_pkcs1_pem(self):
         '''Saves a PKCS#1 PEM-encoded private key file.
 
-        @param keyfile: a PrivateKey object
         @return: contents of a PEM-encoded file that contains the private key.
         '''
 
diff --git a/tests/test_load_save_keys.py b/tests/test_load_save_keys.py
index 5026afd..fca4241 100644
--- a/tests/test_load_save_keys.py
+++ b/tests/test_load_save_keys.py
@@ -8,6 +8,9 @@ import rsa.key
 B64PRIV_DER = 'MC4CAQACBQDeKYlRAgMBAAECBQDHn4npAgMA/icCAwDfxwIDANcXAgInbwIDAMZt'
 PRIVATE_DER = base64.decodestring(B64PRIV_DER)
 
+B64PUB_DER = 'MAwCBQDeKYlRAgMBAAE='
+PUBLIC_DER = base64.decodestring(B64PUB_DER)
+
 PRIVATE_PEM = '''
 -----BEGIN CONFUSING STUFF-----
 Cruft before the key
@@ -26,6 +29,24 @@ CLEAN_PRIVATE_PEM = '''\
 -----END RSA PRIVATE KEY-----
 ''' % B64PRIV_DER
 
+PUBLIC_PEM = '''
+-----BEGIN CONFUSING STUFF-----
+Cruft before the key
+
+-----BEGIN RSA PUBLIC KEY-----
+%s
+-----END RSA PUBLIC KEY-----
+
+Stuff after the key
+-----END CONFUSING STUFF-----
+''' % B64PUB_DER
+
+CLEAN_PUBLIC_PEM = '''\
+-----BEGIN RSA PUBLIC KEY-----
+%s
+-----END RSA PUBLIC KEY-----
+''' % B64PUB_DER
+
 
 class DerTest(unittest.TestCase):
     '''Test saving and loading DER keys.'''
@@ -46,6 +67,22 @@ class DerTest(unittest.TestCase):
 
         self.assertEqual(PRIVATE_DER, der)
 
+    def test_load_public_key(self):
+        '''Test loading public DER keys.'''
+
+        key = rsa.key.PublicKey.load_pkcs1_der(PUBLIC_DER)
+        expected = rsa.key.PublicKey(3727264081, 65537)
+
+        self.assertEqual(expected, key)
+
+    def test_save_public_key(self):
+        '''Test saving public DER keys.'''
+
+        key = rsa.key.PublicKey(3727264081, 65537)
+        der = key.save_pkcs1_der()
+
+        self.assertEqual(PUBLIC_DER, der)
+
 class PemTest(unittest.TestCase):
     '''Test saving and loading PEM keys.'''
 
@@ -66,3 +103,20 @@ class PemTest(unittest.TestCase):
 
         self.assertEqual(CLEAN_PRIVATE_PEM, pem)
 
+    def test_load_public_key(self):
+        '''Test loading public PEM files.'''
+
+        key = rsa.key.PublicKey.load_pkcs1_pem(PUBLIC_PEM)
+        expected = rsa.key.PublicKey(3727264081, 65537)
+
+        self.assertEqual(expected, key)
+
+    def test_save_public_key(self):
+        '''Test saving public PEM files.'''
+
+        key = rsa.key.PublicKey(3727264081, 65537)
+        pem = key.save_pkcs1_pem()
+
+        self.assertEqual(CLEAN_PUBLIC_PEM, pem)
+
+
-- 
cgit v1.2.1