summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorYuval Adam <_@yuv.al>2018-02-19 18:21:53 +0100
committerYuval Adam <_@yuv.al>2018-02-19 18:21:53 +0100
commit975f0b5eb59eedb0f3855e929f867046bec24bf0 (patch)
treed007f45651a053079539f5f45bab955947a97d1f
parent1f9757df41340e142446b1762cb6149f6d4f7223 (diff)
Wrap BinAsciiError
-rw-r--r--tests/test_core.py4
-rw-r--r--viewstate/__main__.py1
-rw-r--r--viewstate/core.py6
3 files changed, 10 insertions, 1 deletions
diff --git a/tests/test_core.py b/tests/test_core.py
index 6d4a8d5..d001d22 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -14,6 +14,10 @@ class TestViewState(object):
vs = ViewState(f.read())
assert vs.is_valid() is True
+ def test_invalid_base64(self):
+ with pytest.raises(ViewStateException):
+ vs = ViewState('hello')
+
def test_invalid_decode(self):
with pytest.raises(ViewStateException):
vs = ViewState()
diff --git a/viewstate/__main__.py b/viewstate/__main__.py
index aa2e25d..298aa14 100644
--- a/viewstate/__main__.py
+++ b/viewstate/__main__.py
@@ -6,6 +6,7 @@ def main():
s = sys.stdin.read()
vs = ViewState(s)
print('Got: ' + s)
+ print(vs.raw)
print(vs.decode())
if __name__ == '__main__':
diff --git a/viewstate/core.py b/viewstate/core.py
index 347ee59..e42dbf1 100644
--- a/viewstate/core.py
+++ b/viewstate/core.py
@@ -1,4 +1,5 @@
from base64 import b64decode, b64encode
+from binascii import Error as BinAsciiError
from .exceptions import ViewStateException
@@ -24,8 +25,11 @@ class ViewState(object):
def __init__(self, base64=''):
self.base64 = base64
- self.raw = b64decode(self.base64)
self.decoded = None
+ try:
+ self.raw = b64decode(self.base64)
+ except BinAsciiError as bae:
+ raise ViewStateException('Cannot decode base64 input')
@property
def preamble(self):