From 549795396cc1810b3be9f2e603bc3498606fbdf9 Mon Sep 17 00:00:00 2001
From: Yuval Adam <_@yuv.al>
Date: Thu, 12 Apr 2018 11:57:00 +0300
Subject: Support raw bytes in CLI

---
 README.rst            |  6 ++++++
 viewstate/__main__.py | 15 +++++++++++----
 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)
-- 
cgit v1.3.1