Support raw bytes in CLI

2 files changed, 17 insertions, 4 deletions

diff --git a/README.rst b/README.rst
index a9573fc..46aaff5 100644
--- a/README.rst
+++ b/README.rst
@@ -42,6 +42,12 @@ Alternatively, the library can be used via command line by directly executing th
 
 Which will pretty-print the decoded data structure.
 
+The command line usage can also accept raw bytes with the `-r` flag:
+
+.. code-block:: shell
+
+  $ cat data.base64 | base64 -d | python -m viewstate -r
+
 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:: python
diff --git a/viewstate/__main__.py b/viewstate/__main__.py
index 558f100..e600cb4 100644
--- a/viewstate/__main__.py
+++ b/viewstate/__main__.py
@@ -3,11 +3,18 @@ import sys
 
 from .viewstate import ViewState
 
-def main():
-    s = sys.stdin.read()
-    vs = ViewState(s)
+
+def main(raw=False):
+    if raw:
+        s = sys.stdin.buffer.read()
+        vs = ViewState(raw=s)
+    else:
+        s = sys.stdin.read()
+        vs = ViewState(s)
     pp = pprint.PrettyPrinter(indent=4)
     pp.pprint(vs.decode())
 
+
 if __name__ == '__main__':
-    main()
+    raw = len(sys.argv) > 1 and sys.argv[1] == '-r'
+    main(raw)