Fix sybrenstuvel#12 Allow pickling of keys.

sybrenstuvel · sybrenstuvel · commit 4bc9733b78cd · 2016-01-22T15:41:40.000+01:00
Pickling is now possible, with the added note that one should never
unpickle from an untrusted or unauthenticated source.
diff --git a/doc/reference.rst b/doc/reference.rst
@@ -21,6 +21,14 @@ Functions
 Classes
 --------------------------------------------------
 
+.. note::
+
+    Storing public and private keys via the `pickle` module is possible.
+    However, it is insecure to load a key from an untrusted source.
+    The pickle module is not secure against erroneous or maliciously
+    constructed data. Never unpickle data received from an untrusted
+    or unauthenticated source.
+
 .. autoclass:: rsa.PublicKey
     :members:
     :inherited-members:
diff --git a/rsa/key.py b/rsa/key.py
@@ -23,6 +23,14 @@
 late as possible, such that other functionality will remain working in absence
 of pyasn1.
 
+.. note::
+
+    Storing public and private keys via the `pickle` module is possible.
+    However, it is insecure to load a key from an untrusted source.
+    The pickle module is not secure against erroneous or maliciously
+    constructed data. Never unpickle data received from an untrusted
+    or unauthenticated source.
+
 """
 
 import logging
@@ -154,6 +162,14 @@ def __getitem__(self, key):
     def __repr__(self):
         return 'PublicKey(%i, %i)' % (self.n, self.e)
 
+    def __getstate__(self):
+        """Returns the key as tuple for pickling."""
+        return self.n, self.e
+
+    def __setstate__(self, state):
+        """Sets the key from tuple."""
+        self.n, self.e = state
+
     def __eq__(self, other):
         if other is None:
             return False
@@ -337,6 +353,14 @@ def __getitem__(self, key):
     def __repr__(self):
         return 'PrivateKey(%(n)i, %(e)i, %(d)i, %(p)i, %(q)i)' % self
 
+    def __getstate__(self):
+        """Returns the key as tuple for pickling."""
+        return self.n, self.e, self.d, self.p, self.q, self.exp1, self.exp2, self.coef
+
+    def __setstate__(self, state):
+        """Sets the key from tuple."""
+        self.n, self.e, self.d, self.p, self.q, self.exp1, self.exp2, self.coef = state
+
     def __eq__(self, other):
         if other is None:
             return False
diff --git a/tests/test_load_save_keys.py b/tests/test_load_save_keys.py
@@ -19,6 +19,7 @@
 import base64
 import unittest
 import os.path
+import pickle
 
 from rsa._compat import b
 
@@ -152,3 +153,22 @@ def test_load_from_disk(self):
 
         self.assertEqual(15945948582725241569, privkey.p)
         self.assertEqual(14617195220284816877, privkey.q)
+
+
+class PickleTest(unittest.TestCase):
+    """Test saving and loading keys by pickling."""
+
+    def test_private_key(self):
+        pk = rsa.key.PrivateKey(3727264081, 65537, 3349121513, 65063, 57287)
+
+        pickled = pickle.dumps(pk)
+        unpickled = pickle.loads(pickled)
+        self.assertEqual(pk, unpickled)
+
+    def test_public_key(self):
+        pk = rsa.key.PublicKey(3727264081, 65537)
+
+        pickled = pickle.dumps(pk)
+        unpickled = pickle.loads(pickled)
+
+        self.assertEqual(pk, unpickled)
