diff options
author | Alex Gaynor <alex.gaynor@gmail.com> | 2019-07-28 22:58:04 -0400 |
---|---|---|
committer | Paul Kehrer <paul.l.kehrer@gmail.com> | 2019-07-28 21:58:04 -0500 |
commit | 9cd41ac714d9bff819ece6d8cdcde064d403c671 (patch) | |
tree | 4ed2502ced1db85417fbf6fe214f59a1525893ff /src/cryptography/hazmat | |
parent | 2c83570f6310cb36553af274eb41dd8e2b96b58e (diff) | |
download | cryptography-9cd41ac714d9bff819ece6d8cdcde064d403c671.tar.gz cryptography-9cd41ac714d9bff819ece6d8cdcde064d403c671.tar.bz2 cryptography-9cd41ac714d9bff819ece6d8cdcde064d403c671.zip |
Make DER reader into a context manager (#4957)
* Make DER reader into a context manager
* Added another test case
* flake8
Diffstat (limited to 'src/cryptography/hazmat')
-rw-r--r-- | src/cryptography/hazmat/_der.py | 12 | ||||
-rw-r--r-- | src/cryptography/hazmat/primitives/asymmetric/utils.py | 9 |
2 files changed, 13 insertions, 8 deletions
diff --git a/src/cryptography/hazmat/_der.py b/src/cryptography/hazmat/_der.py index 3a121a85..51518d64 100644 --- a/src/cryptography/hazmat/_der.py +++ b/src/cryptography/hazmat/_der.py @@ -36,6 +36,13 @@ class DERReader(object): def __init__(self, data): self.data = memoryview(data) + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, tb): + if exc_value is None: + self.check_empty() + def is_empty(self): return len(self.data) == 0 @@ -100,9 +107,8 @@ class DERReader(object): return body def read_single_element(self, expected_tag): - ret = self.read_element(expected_tag) - self.check_empty() - return ret + with self: + return self.read_element(expected_tag) def read_optional_element(self, expected_tag): if len(self.data) > 0 and six.indexbytes(self.data, 0) == expected_tag: diff --git a/src/cryptography/hazmat/primitives/asymmetric/utils.py b/src/cryptography/hazmat/primitives/asymmetric/utils.py index 43d5b9bf..14d2abee 100644 --- a/src/cryptography/hazmat/primitives/asymmetric/utils.py +++ b/src/cryptography/hazmat/primitives/asymmetric/utils.py @@ -12,11 +12,10 @@ from cryptography.hazmat.primitives import hashes def decode_dss_signature(signature): - seq = DERReader(signature).read_single_element(SEQUENCE) - r = seq.read_element(INTEGER).as_integer() - s = seq.read_element(INTEGER).as_integer() - seq.check_empty() - return r, s + with DERReader(signature).read_single_element(SEQUENCE) as seq: + r = seq.read_element(INTEGER).as_integer() + s = seq.read_element(INTEGER).as_integer() + return r, s def encode_dss_signature(r, s): |