From 65d9a80b72215d3b962146302ba847b045331c87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Piotr=20Ma=C5=9Blanka?= <pmaslanka@smok.co> Date: Thu, 14 Oct 2021 20:26:14 +0200 Subject: [PATCH] bugfix --- CHANGELOG.md | 6 ++---- satella/__init__.py | 2 +- satella/parsing.py | 38 +++++++++++++++++++++++++++----------- tests/test_parsing.py | 5 +++++ 4 files changed, 35 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1ee7a1a2..6e182fdf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,3 @@ -# v2.18.1 +# v2.18.2 -* added `BinaryParser.get_remaining_bytes` -* added `BinaryParser.get_remaining_bytes_count` -* added `BinaryParser.get_parser` +* bugfix diff --git a/satella/__init__.py b/satella/__init__.py index 75ab0538..06287ba5 100644 --- a/satella/__init__.py +++ b/satella/__init__.py @@ -1 +1 @@ -__version__ = '2.18.1' +__version__ = '2.18.2' diff --git a/satella/parsing.py b/satella/parsing.py index e98be494..b9cac307 100644 --- a/satella/parsing.py +++ b/satella/parsing.py @@ -27,7 +27,7 @@ class BinaryParser: """ Return the amount of bytes remaining. This will not advance the pointer """ - return self.length - self.pointer + return self.length - self.pointer + self.init_ofs def __init__(self, b_stream: tp.Union[bytes, bytearray], offset: int = 0, length: tp.Optional[int] = None): @@ -41,6 +41,26 @@ class BinaryParser: def __bytes__(self) -> bytes: return self.b_stream[self.init_ofs:self.init_ofs+self.length] + def skip(self, n: int) -> None: + """ + Advance the pointer by n bytes + + :param n: bytes to advance + :raises NotEnoughBytes: not enough bytes remain in the stream! + """ + self.assert_has_bytes(n) + self.pointer += n + + def assert_has_bytes(self, n: int) -> None: + """ + Assert that we have at least n bytes to consume + + :param n: amount of bytes to consume + :raises NotEnoughBytes: not enough bytes remain in the stream! + """ + if self.length + self.init_ofs < self.pointer + n: + raise NotEnoughBytes('Not enough bytes') + def get_parser(self, length: int) -> 'BinaryParser': """ Return a subclassed binary parser providing a window to another binary parser's data. @@ -49,9 +69,9 @@ class BinaryParser: :param length: amount of bytes to view :return: a BinaryParser + :raises NotEnoughBytes: not enough bytes remain in the stream! """ - if self.length < self.pointer + length: - raise NotEnoughBytes('Not enough bytes') + self.assert_has_bytes(length) try: return BinaryParser(self.b_stream, self.pointer, length) finally: @@ -85,8 +105,7 @@ class BinaryParser: :return: bytes returned :raises NotEnoughBytes: not enough bytes remain in the stream! """ - if self.length < self.pointer + n: - raise NotEnoughBytes('Not enough bytes') + self.assert_has_bytes(n) try: return self.b_stream[self.pointer:self.pointer+n] finally: @@ -115,9 +134,7 @@ class BinaryParser: 'get_structs for multiples!' st_len = st.size - if self.length < self.pointer + st_len: - raise NotEnoughBytes('Not enough bytes') - + self.assert_has_bytes(st_len) try: return st.unpack(self.b_stream[self.pointer:self.pointer+st_len])[0] finally: @@ -136,8 +153,7 @@ class BinaryParser: """ st = self._to_struct(st) st_len = st.size - if self.length < self.pointer + st_len: - raise NotEnoughBytes('Not enough bytes') + self.assert_has_bytes(st_len) try: return st.unpack(self.b_stream[self.pointer:self.pointer+st_len]) finally: @@ -149,4 +165,4 @@ class BinaryParser: This will not advance the pointer """ - return self.b_stream[self.pointer:] + return self.b_stream[self.pointer:self.pointer+self.length] diff --git a/tests/test_parsing.py b/tests/test_parsing.py index aa375e27..202cacf9 100644 --- a/tests/test_parsing.py +++ b/tests/test_parsing.py @@ -23,3 +23,8 @@ class TestParsing(unittest.TestCase): self.assertEqual(bp2.get_struct('>L'), 258) self.assertRaises(NotEnoughBytes, lambda: bp2.get_struct('>L')) self.assertEqual(bp.pointer, 4) + bp.reset() + bp.skip(4) + bp3 = bp.get_parser(4) + self.assertEqual(bytes(bp3), b'\x00\x00\x00\xFF') + self.assertRaises(NotEnoughBytes, lambda: bp2.get_parser(4)) -- GitLab