Call in

sybrenstuvel · Feb 24, 2021 · c537db3 · c537db3
1 parent 758562f
commit c537db3
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/rsa/key.py b/rsa/key.py
@@ -251,6 +251,7 @@ def __getstate__(self) -> typing.Tuple[int, int]:
     def __setstate__(self, state: typing.Tuple[int, int]) -> None:
         """Sets the key from tuple."""
         self.n, self.e = state
+        AbstractKey.__init__(self, self.n, self.e)
 
     def __eq__(self, other: typing.Any) -> bool:
         if other is None:
@@ -426,6 +427,7 @@ def __getstate__(self) -> typing.Tuple[int, int, int, int, int, int, int, int]:
     def __setstate__(self, state: typing.Tuple[int, int, int, int, int, int, int, int]) -> None:
         """Sets the key from tuple."""
         self.n, self.e, self.d, self.p, self.q, self.exp1, self.exp2, self.coef = state
+        AbstractKey.__init__(self, self.n, self.e)
 
     def __eq__(self, other: typing.Any) -> bool:
         if other is None:

diff --git a/tests/test_load_save_keys.py b/tests/test_load_save_keys.py
@@ -203,10 +203,15 @@ def test_private_key(self):
         unpickled = pickle.loads(pickled)
         self.assertEqual(pk, unpickled)
 
+        for attr in rsa.key.AbstractKey.__slots__:
+            self.assertTrue(hasattr(unpickled, attr))
+
     def test_public_key(self):
         pk = rsa.key.PublicKey(3727264081, 65537)
 
         pickled = pickle.dumps(pk)
         unpickled = pickle.loads(pickled)
 
         self.assertEqual(pk, unpickled)
+        for attr in rsa.key.AbstractKey.__slots__:
+            self.assertTrue(hasattr(unpickled, attr))