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