diff --git a/CHANGELOG.md b/CHANGELOG.md index 1ee7a1a2bb61a5da92dac21557065929a8a9481a..6e182fdf1fbdc6112ea432dc7b6ce6f705f81361 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 75ab0538ef42bd448949b10e61a73bf3c165f9e0..06287ba5d52ab77c86f74be26b24ee885966ace0 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 e98be49407b758c8e84581ac3f290ecfde4d9995..b9cac307394d716eaa5649f6f1e968ed0a6f49a7 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 aa375e27e0fd3583dabe329c27675b379fc8fc6f..202cacf91c8f013a37d51e01c96aad64128d417e 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))