From 4bc9733b78cd115a742b9486ab11ccbdcb9ca001 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Fri, 22 Jan 2016 15:41:40 +0100
Subject: Fix #12 Allow pickling of keys.

Pickling is now possible, with the added note that one should never
unpickle from an untrusted or unauthenticated source.
---
 doc/reference.rst            |  8 ++++++++
 rsa/key.py                   | 24 ++++++++++++++++++++++++
 tests/test_load_save_keys.py | 20 ++++++++++++++++++++
 3 files changed, 52 insertions(+)

diff --git a/doc/reference.rst b/doc/reference.rst
index d80416a..ce9c1b9 100644
--- 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
index c70db9a..6014709 100644
--- a/rsa/key.py
+++ b/rsa/key.py
@@ -23,6 +23,14 @@ 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.
 
+.. 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 @@ class PublicKey(AbstractKey):
     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 @@ class PrivateKey(AbstractKey):
     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
index 8fc3b22..5ae1596 100644
--- 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 @@ class PemTest(unittest.TestCase):
 
         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)
-- 
cgit v1.2.1