"""The module provides functions to pack and unpack petlib Bn, EcGroup, and EcPt
strucures.
Example:
>>> # Define a custom class, encoder and decoder
>>> class CustomType:
... def __eq__(self, other):
... return isinstance(other, CustomType)
>>>
>>> def enc_custom(obj):
... return b''
>>>
>>> def dec_custom(data):
... return CustomType()
>>>
>>> register_coders(CustomType, 10, enc_custom, dec_custom)
>>>
>>> # Define a structure
>>> G = EcGroup()
>>> custom_obj = CustomType()
>>> test_data = [G, G.generator(), G.order(), custom_obj]
>>>
>>> # Encode and decode custom structure
>>> packed = encode(test_data)
>>> x = decode(packed)
>>> assert x == test_data
"""
try:
import msgpack
except:
print("Generate Documentation")
from .ec import EcGroup, EcPt
from .bn import Bn
__all__ = ["encode", "decode", "register_coders"]
_pack_reg = {}
_unpack_reg = {}
def register_coders(cls, num, enc_func, dec_func):
""" Register a new type for encoding and decoding.
Take a class type, a number, an encoding and a decoding function."""
if num in _unpack_reg or cls in _pack_reg:
raise Exception("Class or number already in use.")
coders = (cls, num, enc_func, dec_func)
_pack_reg[cls] = coders
_unpack_reg[num] = coders
def bn_enc(obj):
if obj < 0:
neg = b"-"
data = (-obj).binary()
else:
neg = b"+"
data = obj.binary()
return neg + data
def bn_dec(data):
num = Bn.from_binary(data[1:])
# Accomodate both Python 2 and Python 3
if data[0] == ord("-") or data[0] == "-":
return -num
return num
def ecg_enc(obj):
# Serialize EcGroup objects
nid = obj.nid()
packed_nid = msgpack.packb(nid)
return packed_nid
def ecg_dec(data):
# Decode EcGroup
nid = msgpack.unpackb(data)
return EcGroup(nid)
def ecpt_enc(obj):
# Serialize EcPt objects
nid = obj.group.nid()
data = obj.export()
packed_data = msgpack.packb((nid, data))
return packed_data
def ecpt_dec(data):
# Decode EcPt
nid, ptdata = msgpack.unpackb(data)
return EcPt.from_binary(ptdata, EcGroup(nid))
def _init_coders():
global _pack_reg, _unpack_reg
_pack_reg, _unpack_reg = {}, {}
register_coders(Bn, 0, bn_enc, bn_dec)
register_coders(EcGroup, 1, ecg_enc, ecg_dec)
register_coders(EcPt, 2, ecpt_enc, ecpt_dec)
# Register default coders
try:
_init_coders()
except:
print("Generate documentation")
def default(obj):
# Serialize Bn objects
for T in _pack_reg:
if isinstance(obj, T):
_, num, enc, _ = _pack_reg[T]
return msgpack.ExtType(num, enc(obj))
raise TypeError("Unknown type: %r" % (type(obj),))
def make_encoder(out_encoder=None):
if out_encoder is None:
return default
else:
def new_encoder(obj):
try:
encoded = default(obj)
return encoded
except BaseException:
return out_encoder(obj)
return new_encoder
def ext_hook(code, data):
if code in _unpack_reg:
_, _, _, dec = _unpack_reg[code]
return dec(data)
# Other
return msgpack.ExtType(code, data)
def make_decoder(custom_decoder=None):
if custom_decoder is None:
return ext_hook
else:
def new_decoder(code, data):
out = ext_hook(code, data)
if not isinstance(out, msgpack.ExtType):
return out
else:
return custom_decoder(code, data)
return new_decoder
def encode(structure, custom_encoder=None):
""" Encode a structure containing petlib objects to a binary format. May define a custom encoder for user classes. """
encoder = make_encoder(custom_encoder)
packed_data = msgpack.packb(structure, default=encoder, use_bin_type=True)
return packed_data
def decode(packed_data, custom_decoder=None):
""" Decode a binary byte sequence into a structure containing pelib objects. May define a custom decoder for custom classes. """
decoder = make_decoder(custom_decoder)
structure = msgpack.unpackb(
packed_data,
ext_hook=decoder,
encoding='utf-8')
return structure
# --- TESTS ---
def test_basic():
x = [b'spam', u'egg']
packed = msgpack.packb(x, use_bin_type=True)
y = msgpack.unpackb(packed, encoding='utf-8')
assert x == y
def test_bn():
bn1, bn2 = Bn(1), Bn(2)
test_data = [bn1, bn2, -bn1, -bn2]
packed = msgpack.packb(test_data, default=default, use_bin_type=True)
x = msgpack.unpackb(packed, ext_hook=ext_hook, encoding='utf-8')
assert x == test_data
def test_ecgroup():
G = EcGroup()
test_data = [G]
packed = msgpack.packb(test_data, default=default, use_bin_type=True)
x = msgpack.unpackb(packed, ext_hook=ext_hook, encoding='utf-8')
assert x == test_data
def test_ecpt():
G = EcGroup()
test_data = [G.generator()]
packed = msgpack.packb(test_data, default=default, use_bin_type=True)
x = msgpack.unpackb(packed, ext_hook=ext_hook, encoding='utf-8')
assert x == test_data
def test_mixed():
G = EcGroup()
test_data = [G, G.generator(), G.order()]
packed = msgpack.packb(test_data, default=default, use_bin_type=True)
x = msgpack.unpackb(packed, ext_hook=ext_hook, encoding='utf-8')
assert x == test_data
def test_enc_dec():
G = EcGroup()
test_data = [G, G.generator(), G.order()]
packed = encode(test_data)
x = decode(packed)
assert x == test_data
def test_enc_dec_dict():
G = EcGroup()
# , G.generator():"1", "2":G.order()}
test_data = {G.order(): [G, G.generator()]}
packed = encode(test_data)
x = decode(packed)
assert x[G.order()] == test_data[G.order()]
def test_enc_dec_custom():
# Define a custom class, encoder and decoder
class CustomClass:
def __eq__(self, other):
return isinstance(other, CustomClass)
def enc_CustomClass(obj):
if isinstance(obj, CustomClass):
return msgpack.ExtType(11, b'')
raise TypeError("Unknown type: %r" % (obj,))
def dec_CustomClass(code, data):
if code == 11:
return CustomClass()
return msgpack.ExtType(code, data)
# Define a structure
G = EcGroup()
custom_obj = CustomClass()
test_data = [G, G.generator(), G.order(), custom_obj]
# Encode and decode custom structure
packed = encode(test_data, enc_CustomClass)
x = decode(packed, dec_CustomClass)
assert x == test_data
def test_streaming():
# Define a custom class, encoder and decoder
class CustomClass2:
def __eq__(self, other):
return isinstance(other, CustomClass2)
def enc_CustomClass(obj):
if isinstance(obj, CustomClass2):
return msgpack.ExtType(12, b'')
raise TypeError("Unknown type: %r" % (obj,))
def dec_CustomClass(code, data):
if code == 12:
return CustomClass2()
return msgpack.ExtType(code, data)
# Define a structure
G = EcGroup()
custom_obj = CustomClass2()
test_data = [G, G.generator(), G.order(), custom_obj]
packed1 = encode(test_data, enc_CustomClass)
packed2 = encode(test_data, enc_CustomClass)
data = packed1 + packed2
decoder = make_decoder(dec_CustomClass)
Up = msgpack.Unpacker(ext_hook=decoder)
Up.feed(data)
for o in Up:
assert o == test_data
def test_docstring():
# Define a custom class, encoder and decoder
class CustomType:
def __eq__(self, other):
return isinstance(other, CustomType)
def enc_custom(obj):
return b''
def dec_custom(data):
return CustomType()
_init_coders()
register_coders(CustomType, 14, enc_custom, dec_custom)
assert CustomType in _pack_reg
# Define a structure
G = EcGroup()
custom_obj = CustomType()
test_data = [G, G.generator(), G.order(), custom_obj]
# Encode and decode custom structure
packed = encode(test_data)
x = decode(packed)
assert x == test_data
_init_coders()
