diff options
| author | Yuval Adam <_@yuv.al> | 2018-02-19 18:21:53 +0100 |
|---|---|---|
| committer | Yuval Adam <_@yuv.al> | 2018-02-19 18:21:53 +0100 |
| commit | 975f0b5eb59eedb0f3855e929f867046bec24bf0 (patch) | |
| tree | d007f45651a053079539f5f45bab955947a97d1f | |
| parent | 1f9757df41340e142446b1762cb6149f6d4f7223 (diff) | |
Wrap BinAsciiError
| -rw-r--r-- | tests/test_core.py | 4 | ||||
| -rw-r--r-- | viewstate/__main__.py | 1 | ||||
| -rw-r--r-- | viewstate/core.py | 6 |
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): |
