add encrypted viewstate support. Issue yuvadm#10

viewstate/__main__.py

@@ -1,20 +1,29 @@
 import pprint
 import sys
+import argparse
 
 from .viewstate import ViewState
+from .keymodes import Encryption, MAC
 
 
-def main(raw=False):
-    if raw:
+def main(options):
+    if options.raw:
         s = sys.stdin.buffer.read()
-        vs = ViewState(raw=s)
+        vs = ViewState(raw=s, encryption_mode=options.encryption_mode, encryption_key=bytearray.fromhex(options.encryption_key), mac_mode=options.mac_mode)
     else:
         s = sys.stdin.read()
-        vs = ViewState(s)
+        vs = ViewState(s, encryption_mode=options.encryption_mode, encryption_key=bytearray.fromhex(options.encryption_key), mac_mode=options.mac_mode)
     pp = pprint.PrettyPrinter(indent=4)
     pp.pprint(vs.decode())
 
+def getOptions(args=sys.argv[1:]):
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument("-r", dest='raw', action='store_true', default=False, help="Raw mode")
+    parser.add_argument("-e", dest='encryption_mode', action='store', type=Encryption, choices=list(Encryption), help="Encryption algorithm")
+    parser.add_argument("-k", dest='encryption_key', action='store', help="Encryption key in hexadecimal")
+    parser.add_argument("-m", dest='mac_mode', action='store', type=MAC, choices=list(MAC), help="MAC validation algorithm")
+    return parser.parse_args(args)
 
 if __name__ == "__main__":
-    raw = len(sys.argv) > 1 and sys.argv[1] == "-r"
-    main(raw)
+    options = getOptions()
+    main(options)

viewstate/keymodes.py

@@ -0,0 +1,19 @@
+from enum import Enum
+
+class Encryption(Enum):
+    DES = 'des'
+    DES3 = '3des'
+    AES = 'aes'
+
+    def __str__(self):
+        return self.value
+
+class MAC(Enum): #https://referencesource.microsoft.com/#System.Web/Configuration/MachineKeyValidation.cs,a357c185ad2e9c71
+    HMACMD5 = 'hmac_md5'
+    HMACSHA1 = 'hmac_sha1'
+    HMACSHA256 = 'hmac_sha256'
+    HMACSHA384 = 'hmac_sha384'
+    HMACSHA512 = 'hmac_sha512'
+
+    def __str__(self):
+        return self.value

viewstate/viewstate.py

@@ -1,12 +1,17 @@
 from base64 import b64decode, b64encode
 from binascii import Error as BinAsciiError
+try:
+    from Crypto.Cipher import AES, DES3, DES
+except ImportError:
+    from Cryptodome.Cipher import AES, DES3, DES
 
+from .keymodes import Encryption, MAC
 from .exceptions import ViewStateException
 from .parse import Parser
 
 
 class ViewState(object):
-    def __init__(self, base64=None, raw=None):
+    def __init__(self, base64=None, raw=None, encryption_mode=None, encryption_key=None, mac_mode=None):
         if base64:
             self.base64 = base64
             try:
@@ -15,6 +20,19 @@ def __init__(self, base64=None, raw=None):
                 raise ViewStateException("Cannot decode base64 input")
         elif raw:
             self.raw = raw
+        self.decrypt = False
+        if isinstance(encryption_mode, Encryption):
+            if encryption_key is not None and isinstance(mac_mode, MAC):
+                self.encryption_mode = encryption_mode
+                self.encryption_key = encryption_key
+                self.mac_mode = mac_mode
+                self.decrypt = True
+            else:
+                raise ViewStateException("Cannot decrypt without key or mac_mode")
+        else:
+            self.encryption_mode = None
+            self.encryption_key = None
+            self.mac_mode = None
         self.decoded = None
         self.mac = None
         self.signature = None
@@ -38,18 +56,50 @@ def is_valid(self):
             return False
 
     def decode(self):
-        if not self.is_valid():
-            raise ViewStateException("Cannot decode invalid viewstate, bad preamble")
+        if self.decrypt:
+            if self.mac_mode == MAC.HMACMD5:
+                hashlen = 16
+            elif self.mac_mode == MAC.HMACSHA1:
+                hashlen = 20
+            elif self.mac_mode == MAC.HMACSHA256:
+                hashlen = 32
+            elif self.mac_mode == MAC.HMACSHA384:
+                hashlen = 48
+            elif self.mac_mode == MAC.HMACSHA512:
+                hashlen = 64
 
-        self.decoded, self.remainder = Parser.parse(self.body)
+            if self.encryption_mode == Encryption.DES:
+                decryptor = DES.new(self.encryption_key, DES.MODE_CBC, bytearray(8))
+                blockpadlen = 8
+            elif self.encryption_mode == Encryption.DES3:
+                decryptor = DES3.new(self.encryption_key, DES3.MODE_CBC, bytearray(8))
+                blockpadlen = 24
+            elif self.encryption_mode == Encryption.AES:
+                decryptor = AES.new(self.encryption_key, AES.MODE_CBC, bytearray(16))
+                blockpadlen = 32
 
-        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
+            self.signature = self.raw[-hashlen:]
+            self.mac = self.mac_mode.value
+            decrypted = decryptor.decrypt(self.raw[:-hashlen])
+            self.raw = decrypted[blockpadlen:]
+
+            if not self.is_valid():
+                raise ViewStateException('Cannot decode invalid viewstate, bad preamble')
+
+            self.decoded, self.remainder = Parser.parse(self.body)
+        else:
+            if not self.is_valid():
+                raise ViewStateException('Cannot decode invalid viewstate, bad preamble')
+
+            self.decoded, self.remainder = Parser.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