Fix test_rekey_and_reencryption

nkms/crypto/keypairs.py

@@ -27,16 +27,18 @@ class EncryptingKeypair(object):
         symm_key, enc_symm_key = self.pre.encapsulate(pubkey)
         return (symm_key, enc_symm_key)
 
-    def decrypt_key(self, enc_key):
+    def decrypt_key(self, enc_key, privkey=None):
         """
         Decrypts an ECIES encrypted symmetric key.
 
         :param int enc_key: The ECIES encrypted key as an integer
+        :param bytes privkey: The privkey to decapsulate from
 
         :rtype: int
         :return: Decrypted key as an integer
         """
-        return self.pre.decapsulate(self.priv_key, enc_key)
+        priv_key = privkey or self.priv_key
+        return self.pre.decapsulate(priv_key, enc_key)
 
     def rekey(self, privkey_a, privkey_b):
         """

nkms/crypto/keyring.py

@@ -1,4 +1,5 @@
 import sha3
+import npre.elliptic_curve as ec
 from nacl.utils import random
 from nacl.secret import SecretBox
 from nkms.crypto.keypairs import SigningKeypair, EncryptingKeypair
@@ -74,35 +75,39 @@ class KeyRing(object):
         symm_key, enc_symm_key = self.enc_keypair.generate_key()
         return (symm_key, enc_symm_key)
 
-    def decrypt_key(self, enc_key):
+    def decrypt_key(self, enc_key, privkey=None):
         """
         Decrypts an ECIES encrypted symmetric key.
 
         :param EncryptedKey enc_key: ECIES encrypted key in bytes
+        :param bytes privkey: The privkey to decrypt with
 
         :rtype: bytes
         :return: Bytestring of the decrypted symmetric key
         """
-        return self.enc_keypair.decrypt_key(enc_key)
+        return self.enc_keypair.decrypt_key(enc_key, privkey)
 
-    def rekey(self, privkey_a, pubkey_b, symm_key):
+    def rekey(self, privkey_a, pubkey_b):
         """
         Generates a re-encryption key in interactive mode.
 
         :param bytes privkey_a: Alive's private key
         :param bytes pubkey_b: Bob's public key
-        :param bytes symm_key: Symmetric key to encrypt the ephemeral key
 
         :rtype: bytes
         :return: Bytestring of a re-encryption key
         """
+        # Generate an ephemeral keypair
         priv_e = self.enc_keypair.pre.gen_priv()
-        priv_e_bytes = int(priv_e).to_bytes(self.enc_keypair.KEYSIZE,
-                                            byteorder='big')
+        priv_e_bytes = ec.serialize(priv_e)[1:]
+
+        # Encrypt ephemeral key with an ECIES generated key
+        symm_key_bob, enc_symm_key_bob = self.enc_keypair.generate_key(
+                                                            pubkey=pubkey_b)
+        enc_priv_e = self.symm_encrypt(symm_key_bob, priv_e_bytes)
 
         reenc_key = self.enc_keypair.rekey(self.enc_privkey, priv_e)
-        enc_priv_e = self.symm_encrypt(symm_key, priv_e_bytes)
-        return (reenc_key, enc_priv_e)
+        return (reenc_key, enc_symm_key_bob, enc_priv_e)
 
     def reencrypt(self, reenc_key, ciphertext):
         """

tests/crypto/test_keyring.py

@@ -1,7 +1,9 @@
 import unittest
 import msgpack
 import random
+import npre.elliptic_curve as ec
 from nkms.crypto.keyring import KeyRing
+from nacl.secret import SecretBox
 
 
 class TestKeyRing(unittest.TestCase):
@@ -46,16 +48,31 @@ class TestKeyRing(unittest.TestCase):
         self.assertTrue(raw_key == dec_key)
 
     def test_rekey_and_reencryption(self):
-        raw_key, enc_key = self.keyring_a.generate_key()
-        reenc_key, enc_priv_e = self.keyring_a.rekey(self.keyring_a.enc_privkey,
-                                                     self.keyring_b.enc_pubkey,
-                                                     raw_key)
+        # Generate random key for Alice's data
+        symm_key_alice, enc_symm_key_alice = self.keyring_a.generate_key()
 
-        rekey_enc_key = self.keyring_a.reencrypt(reenc_key, enc_key)
+        # Generate the re-encryption key and the ephemeral key data
+        reenc_key, enc_symm_key_bob, enc_priv_e = self.keyring_a.rekey(
+                                                    self.keyring_a.enc_privkey,
+                                                    self.keyring_b.enc_pubkey)
 
-        dec_key = self.keyring_b.decrypt_key(rekey_enc_key)
-        self.assertEqual(32, len(dec_key))
-        self.assertTrue(dec_key == raw_key)
+        # Re-encrypt Alice's symm key
+        enc_symm_key_ab = self.keyring_a.reencrypt(reenc_key,
+                                                   enc_symm_key_alice)
+
+        # Decapsulate the ECIES encrypted key
+        dec_symm_key_bob = self.keyring_b.decrypt_key(enc_symm_key_bob)
+        self.assertEqual(32, len(dec_symm_key_bob))
+
+        # Decrypt the ephemeral private key
+        dec_priv_e = self.keyring_b.symm_decrypt(dec_symm_key_bob, enc_priv_e)
+        self.assertEqual(32, len(dec_priv_e))
+        dec_priv_e = int.from_bytes(dec_priv_e, byteorder='big')
+
+        dec_key = self.keyring_b.decrypt_key(enc_symm_key_ab,
+                                             privkey=dec_priv_e)
+
+        self.assertEqual(dec_key, symm_key_alice)
 
     def test_split_key_sharing(self):
         raw_key, enc_key = self.keyring_a.generate_key()