diff --git a/npre/umbral.py b/npre/umbral.py index 5a3f542..0e5dd52 100644 --- a/npre/umbral.py +++ b/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,7 +82,7 @@ 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 @@ -90,7 +90,7 @@ def split_rekey(self, priv_a, priv_b, threshold, N): 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) diff --git a/tests/test_umbral.py b/tests/test_umbral.py index 0d4b613..5d4dbc6 100644 --- a/tests/test_umbral.py +++ b/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 \ No newline at end of file