From 27698f929a4f71d84e508b51727699f9f35f3dc0 Mon Sep 17 00:00:00 2001
From: Kevin Froman <beardog@mailbox.org>
Date: Thu, 10 Dec 2020 01:30:26 +0000
Subject: [PATCH] added getter for unserialized metadata

---
 kasten/generator/__init__.py             | 14 +++++++++---
 kasten/main.py                           | 13 ++++++++++-
 setup.py                                 |  2 +-
 tests/test_kasten_with_base_generator.py | 28 ++++++++++++++++++++++--
 4 files changed, 50 insertions(+), 7 deletions(-)

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()