From 8638d1c254e7813de469e86eaa290b413c547ff5 Mon Sep 17 00:00:00 2001
From: Yuval Adam <_@yuv.al>
Date: Tue, 20 Feb 2018 15:01:53 +0100
Subject: Better int parsing

---
 tests/test_core.py    |  4 ++++
 viewstate/__main__.py |  2 +-
 viewstate/parse.py    | 27 ++++++++++++++++++++++++++-
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/tests/test_core.py b/tests/test_core.py
index 2a3b11c..2a0078a 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -23,6 +23,10 @@ class TestViewState(object):
             vs = ViewState(raw=b'\x01\x02')
             vs.decode()
 
+    def test_parse_int_value(self):
+        vs = ViewState(raw=b'\xff\x01\x02\x88\x01')
+        assert vs.decode() == 136
+
     def test_parse_const_value(self):
         vs = ViewState(raw=b'\xff\x01\x67')
         assert vs.decode() is True
diff --git a/viewstate/__main__.py b/viewstate/__main__.py
index 16f2919..558f100 100644
--- a/viewstate/__main__.py
+++ b/viewstate/__main__.py
@@ -1,7 +1,7 @@
 import pprint
 import sys
 
-from .core import ViewState
+from .viewstate import ViewState
 
 def main():
     s = sys.stdin.read()
diff --git a/viewstate/parse.py b/viewstate/parse.py
index 0692c5c..7620f25 100644
--- a/viewstate/parse.py
+++ b/viewstate/parse.py
@@ -11,11 +11,32 @@ CONSTS = {
 def parse_const(b):
     return CONSTS.get(b, None)
 
+def parse_int(b):
+    n = 0
+    bits = 0
+    i = 0
+    while (bits < 32):
+        tmp = b[i]
+        i += 1
+        n |= (tmp & 0x7f) << bits
+        if not tmp & 0x80:
+            return n, b[i:]
+        bits += 7
+    return n, b[i:]  # overflow
+
 def parse_string(b):
     n = b[0]
     s = b[1:n+1]
     return s.decode(), b[n+1:]
 
+def parse_enum(b):
+    if b[0] == 0x29:
+        n, remain = parse_int(b[1:])
+        for _ in range(n):
+            return remain[:n], remain[n+1:]
+    else:
+        pass
+
 def parse_pair(b):
     first, remain = parse(b)
     second, remain = parse(remain)
@@ -48,8 +69,12 @@ def parse(b):
 
     if 100 <= b[0] <= 104:
         return parse_const(b[0]), b[1:]
-    elif b[0] == 0x5:
+    elif b[0] == 0x2:
+        return parse_int(b[1:])
+    elif b[0] in (0x5, 0x1e):
         return parse_string(b[1:])
+    elif b[0] == 0xb:
+        return parse_enum(b[1:])
     elif b[0] == 0xf:
         return parse_pair(b[1:])
     elif b[0] == 0x16:
-- 
cgit v1.3.1