Add RekeyFrag class

npre/umbral.py

@@ -4,6 +4,7 @@
 
 import npre.elliptic_curve as ec
 from npre import curves
+from npre.constants import UNKNOWN_KFRAG
 from typing import Union
 from sha3 import keccak_256 as keccak
 from collections import namedtuple
@@ -12,7 +13,6 @@
 
 
 EncryptedKey = namedtuple('EncryptedKey', ['ekey', 're_id'])
-RekeyFrag = namedtuple('RekeyFrag', ['id', 'key'])
 
 # XXX serialization probably should be done through decorators
 # XXX write tests
@@ -82,15 +82,15 @@ def save_key(self, key):
     def rekey(self, priv1, priv2, dtype=None):
         # Same as in BBS98
         rk = priv1 * (~priv2)
-        return RekeyFrag(id=None, key=rk)
+        return RekeyFrag(id=None, key=rk, pre=self)
 
     def split_rekey(self, priv_a, priv_b, threshold, N):
         coeffs = [priv_a * (~priv_b)]  # Standard rekey
         coeffs += [ec.random(self.ecgroup, ec.ZR) for _ in range(threshold - 1)]
 
         ids = [ec.random(self.ecgroup, ec.ZR) for _ in range(N)]
         rk_shares = [
-                RekeyFrag(id, key=poly_eval(coeffs, id))
+                RekeyFrag(id, key=poly_eval(coeffs, id), pre=self)
                 for id in ids]
 
         return rk_shares
@@ -129,3 +129,30 @@ def decapsulate(self, priv_key, ekey):
         shared_key = ekey.ekey ** priv_key
         key = self.kdf(shared_key)
         return key
+
+
+class RekeyFrag(object):
+
+    _pre = PRE()
+
+    def __init__(self, id, key, pre=None):
+        self.id = id
+        self.key = key
+        if pre is None:
+            pre = self._pre
+        self.pre = pre
+
+    def __bytes__(self):
+        return ec.serialize(self.id) + ec.serialize(self.key)
+
+    def __eq__(self, other_kfrag):
+        if other_kfrag is UNKNOWN_KFRAG:
+            return False
+        return bytes(self) == bytes(other_kfrag)
+
+    @classmethod
+    def from_bytes(cls, kfrag_bytes, pre=None):
+        pre = pre or cls._pre
+        return RekeyFrag(id=ec.deserialize(pre.ecgroup, kfrag_bytes[:len(kfrag_bytes) // 2]),
+                         key=ec.deserialize(pre.ecgroup, kfrag_bytes[len(kfrag_bytes) // 2:]),
+                         pre=pre)

tests/test_umbral.py

@@ -1,5 +1,7 @@
 import pytest
 from npre import umbral
+from npre.umbral import RekeyFrag
+import npre.elliptic_curve as ec
 
 
 def test_encrypt_decrypt():
@@ -60,3 +62,22 @@ def test_m_of_n(N, threshold):
 
     sym_key_2 = pre.decapsulate(priv_bob, ekey_bob)
     assert sym_key_2 == sym_key
+
+    return kfrags, pre
+
+
+def test_kfrag_serialization():
+    kfrags, pre = test_m_of_n(5, 5)
+    some_particular_kfrag = kfrags[3]
+    original_id = some_particular_kfrag.id
+    serialized_id = ec.serialize(original_id)
+    deserialized_id = ec.deserialize(pre.ecgroup, serialized_id)
+    assert deserialized_id == original_id
+
+
+def test_frag_as_bytes():
+    kfrags, pre = test_m_of_n(5, 5)
+    some_particular_kfrag = kfrags[3]
+    kfrag_as_bytes = bytes(some_particular_kfrag)
+    back_to_kfrag = RekeyFrag.from_bytes(kfrag_as_bytes)
+    assert some_particular_kfrag == back_to_kfrag