Extract body property and test for bad preamble decoding

2 files changed, 29 insertions, 1 deletions

diff --git a/test.py b/test.py
index 84249b2..2cda602 100644
--- a/test.py
+++ b/test.py
@@ -9,5 +9,10 @@ class TestViewState(unittest.TestCase):
             vs = ViewState(f.read())
             self.assertTrue(vs.is_valid())
 
+    def test_invalid_decode(self):
+        with self.assertRaises(Exception):
+            vs = ViewState('\x01\x02')
+            vs.decode()
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/viewstate.py b/viewstate.py
index 30da967..d1c9cf8 100644
--- a/viewstate.py
+++ b/viewstate.py
@@ -1,14 +1,37 @@
 from base64 import b64decode, b64encode
 
 
+CONSTS = {
+    '\x64': {},
+    '\x65': '',
+    '\x66': 0,
+    '\x67': True,
+    '\x68': False
+}
+
+def parse():
+    pass
+
+
 class ViewState(object):
 
     def __init__(self, base64_encoded):
         self.base64_encoded = base64_encoded
         self.raw = b64decode(self.base64_encoded)
+        self.decoded = None
+
+    @property
+    def body(self):
+        return self.raw[:2]
 
     def is_valid(self):
         format_marker = b'\xff'
         version_marker = b'\x01'
         preamble = format_marker + version_marker
-        return self.raw[:2] == preamble
+        return self.body == preamble
+
+    def decode(self):
+        if not self.is_valid():
+            raise Exception('Cannot decode invalid viewstate, bad preamble')
+        self.decoded = parse(self.body)
+        return self.decoded