diff --git a/kasten/generator/__init__.py b/kasten/generator/__init__.py index fbbce49..35b5389 100644 --- a/kasten/generator/__init__.py +++ b/kasten/generator/__init__.py @@ -41,11 +41,19 @@ class KastenBaseGenerator: class KastenMimcGenerator: @classmethod - def generate(cls, packed_bytes: KastenPacked, rounds: int = 5000) -> Kasten: - return Kasten(vdf_create(packed_bytes, rounds, dec=True).to_bytes(64, "big"), packed_bytes, cls, auto_check_generator=False) + def generate( + cls, packed_bytes: KastenPacked, rounds: int = 5000) -> Kasten: + return Kasten( + vdf_create( + packed_bytes, + rounds, dec=True + ).to_bytes( + 64, "big"), packed_bytes, cls, auto_check_generator=False) @staticmethod - def validate_id(hash: KastenChecksum, packed_bytes: KastenPacked, rounds = 5000) -> None: + def validate_id( + hash: KastenChecksum, + packed_bytes: KastenPacked, rounds=5000) -> None: try: hash = int.from_bytes(hash, byteorder="big") except TypeError: diff --git a/kasten/main.py b/kasten/main.py index 80d536e..80c931f 100644 --- a/kasten/main.py +++ b/kasten/main.py @@ -32,7 +32,9 @@ class Kasten: self.id = id self.generator = generator header, data = packed_bytes.split(b'\n', 1) - header = unpackb(header, strict_map_key=True) + header = unpackb( + header, + strict_map_key=True) self.header = header self.data = data self.additional_generator_args = list(additional_generator_args) @@ -51,8 +53,15 @@ class Kasten: # Getters are gross, but they are used here to preserve space def get_packed(self) -> KastenPacked: + def _get_or_none(func): + try: + ret = func() + except IndexError: + return None + return ret return pack.pack(self.data, self.get_data_type(), self.get_encryption_mode(), + app_metadata=_get_or_none(self.get_metadata), timestamp=self.get_timestamp()) def get_data_type(self) -> str: return self.header[0] @@ -60,3 +69,5 @@ class Kasten: def get_encryption_mode(self): return self.header[1] def get_timestamp(self): return self.header[2] + + def get_metadata(self): return self.header[3] diff --git a/setup.py b/setup.py index 769c0e0..4c120a1 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ with open(path.join(this_directory, 'README.md'), encoding='utf-8') as f: long_description = f.read() setup(name='kasten', - version='0.1.0', + version='0.2.0', description='Efficient, secure data serialization format with extensibility.', author='Kevin Froman', author_email='beardog@mailbox.org', diff --git a/tests/test_kasten_with_base_generator.py b/tests/test_kasten_with_base_generator.py index 97a0a96..04e244b 100644 --- a/tests/test_kasten_with_base_generator.py +++ b/tests/test_kasten_with_base_generator.py @@ -3,7 +3,7 @@ from hashlib import sha3_384 from kasten import Kasten from kasten import exceptions -from kasten.generator import KastenBaseGenerator +from kasten.generator import KastenBaseGenerator, pack class TestKastenBaseGenerator(unittest.TestCase): @@ -13,7 +13,31 @@ class TestKastenBaseGenerator(unittest.TestCase): def test_kasten_invalid(self): k = b'\x93\xa3txt\x01\xce^\x97\xe3\xdc\ntest' - self.assertRaises(exceptions.InvalidID, Kasten, sha3_384(k + b'invalid').digest(), k, KastenBaseGenerator) + self.assertRaises( + exceptions.InvalidID, + Kasten, sha3_384(k + b'invalid').digest(), k, KastenBaseGenerator) + + def test_kasten_get_metadata(self): + metadata = {"name": "john", "raw": b"are we having fun yet?"} + packed = pack.pack(b"test msg", "tst", 0, app_metadata=metadata) + K = Kasten(sha3_384(packed).digest(), packed, KastenBaseGenerator) + self.assertEqual(K.get_metadata(), metadata) + + def test_kasten_unsafe_deserialize(self): + class Person: + def __init__(self, n): + self.name = n + kevin = Person("kevin") + metadata = { + "name": "john", "raw": b"are we having fun yet?", "person": kevin} + try: + packed = pack.pack(b"test msg", "tst", 0, app_metadata=metadata) + K = Kasten(sha3_384(packed).digest(), packed, KastenBaseGenerator) + except TypeError: + pass + else: + raise Exception("Serialized custom type in packer") + unittest.main()