From ba0b06637bcec3a01d8e6a11884fbad1a0f5d612 Mon Sep 17 00:00:00 2001
From: Yuval Adam <_@yuv.al>
Date: Mon, 19 Mar 2018 11:58:49 +0200
Subject: Add support for HMAC signatures, fixes #4

---
 README.rst              | 25 +++++++++++++++++++------
 setup.py                |  2 +-
 tests/test_viewstate.py | 20 ++++++++++++++++++++
 viewstate/__main__.py   |  1 +
 viewstate/viewstate.py  | 13 +++++++++++++
 5 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/README.rst b/README.rst
index 686bdc3..cb07b58 100644
--- a/README.rst
+++ b/README.rst
@@ -20,17 +20,19 @@ The Viewstate decoder accepts Base64 encoded .NET viewstate data and returns the
 
 There are two main ways to use this package. First, it can be used as an imported library with the following typical use case:
 
-.. code-block:: python
+.. code-block:: pycon
 
-  from viewstate import ViewState
-  vs = ViewState(base64EncodedViewState)
-  decoded_state = vs.decode()
+  >>> from viewstate import ViewState
+  >>> base64_encoded_viewstate = '/wEPBQVhYmNkZQ9nAgE='
+  >>> vs = ViewState(base64_encoded_viewstate)
+  >>> vs.decode()
+  ('abcde', (True, 1))
 
 It is also possible to feed the raw bytes directly:
 
-.. code-block:: python
+.. code-block:: pycon
 
-  vs = ViewState(raw=b'\xff\x01....')
+  >>> vs = ViewState(raw=b'\xff\x01....')
 
 Alternatively, the library can be used via command line by directly executing the module:
 
@@ -40,6 +42,17 @@ Alternatively, the library can be used via command line by directly executing th
 
 Which will pretty-print the decoded data structure.
 
+Viewstate HMAC signatures are also supported. In case there are any remaining bytes after parsing, they are assumed to be HMAC signatures, with the types estimated according to signature length.
+
+.. code-block:: pycon
+
+   >>> vs = ViewState(signed_view_state)
+   >>> vs.decode()
+   >>> vs.mac
+   'hmac_sha256'
+   >>> vs.signature
+   b'....'
+
 Development
 -----------
 
diff --git a/setup.py b/setup.py
index 78cb128..dd2ae42 100644
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,7 @@ setup(
     name='viewstate',
     author='Yuval Adam',
     author_email='_@yuv.al',
-    version='0.3.0',
+    version='0.3.1',
     description='.NET viewstate decoder',
     long_description=long_description,
     url='https://github.com/yuvadm/viewstate',
diff --git a/tests/test_viewstate.py b/tests/test_viewstate.py
index ca20167..9cc9285 100644
--- a/tests/test_viewstate.py
+++ b/tests/test_viewstate.py
@@ -26,3 +26,23 @@ class TestViewState(object):
         with pytest.raises(ViewStateException):
             vs = ViewState(raw=b'\x01\x02')
             vs.decode()
+
+    def test_no_signature(self):
+        vs = ViewState(raw=b'\xff\x01e')
+        vs.decode()
+        assert vs.mac is None
+        assert vs.signature is None
+
+    def test_macs(self):
+        MAC_LENGTHS = {
+            'hmac_sha1': 20,
+            'hmac_sha256': 32,
+            'unknown': 5  # could be any other value
+        }
+
+        for mac, n in MAC_LENGTHS.items():
+            sig = b'\x55' * n  # not a real signature, just testing length
+            vs = ViewState(raw=b'\xff\x01d' + sig)
+            vs.decode()
+            assert vs.mac == mac
+            assert vs.signature == sig
diff --git a/viewstate/__main__.py b/viewstate/__main__.py
index 558f100..217d49f 100644
--- a/viewstate/__main__.py
+++ b/viewstate/__main__.py
@@ -8,6 +8,7 @@ def main():
     vs = ViewState(s)
     pp = pprint.PrettyPrinter(indent=4)
     pp.pprint(vs.decode())
+    print(vs.signed, vs.mac)
 
 if __name__ == '__main__':
     main()
diff --git a/viewstate/viewstate.py b/viewstate/viewstate.py
index 81ef31e..994f1b1 100644
--- a/viewstate/viewstate.py
+++ b/viewstate/viewstate.py
@@ -17,6 +17,8 @@ class ViewState(object):
         elif raw:
             self.raw = raw
         self.decoded = None
+        self.mac = None
+        self.signature = None
 
     @property
     def preamble(self):
@@ -39,5 +41,16 @@ class ViewState(object):
     def decode(self):
         if not self.is_valid():
             raise ViewStateException('Cannot decode invalid viewstate, bad preamble')
+
         self.decoded, self.remainder = parse(self.body)
+
+        if self.remainder:
+            if len(self.remainder) == 20:
+                self.mac = 'hmac_sha1'
+            elif len(self.remainder) == 32:
+                self.mac = 'hmac_sha256'
+            else:
+                self.mac = 'unknown'
+            self.signature = self.remainder
+
         return self.decoded
-- 
cgit v1.3.1