Adding Strict Parsing Mode and RemainChunkTransaction

Author: dannyhann

.coveragerc

@@ -2,6 +2,8 @@
 omit =
     venv/*
     bin/*
+    build/*
+    dist/*
     .pytest_cache/*
     LICENSE
     README.md

coverage.svg

@@ -116,19 +116,28 @@ def to_dict(self):
 @dataclass
 class ParsedResults:
     results: List[ParsedResult]
+    remain_data: str = None
 
     @property
     def has_results(self):
         return len(self.results) > 0
 
+    @property
+    def has_remain_data(self):
+        return self.remain_data is not None
+
     def __getitem__(self, item):
         return self.results[item]
 
     def to_dict(self):
         results = [result.to_dict() for result in self.results]
-        return dict(
+        dict_results = dict(
             results=results,
         )
+        if self.has_remain_data:
+            dict_results["remain_data"] = self.remain_data
+
+        return dict_results
 
 
 class State(Enum):
@@ -197,6 +206,44 @@ def hex_string_to_utf8(cls, string) -> str:
         string = string.replace(" ", "")
         return binascii.unhexlify(string).decode("utf-8")
 
+    @classmethod
+    def chunk_to_hex_string(cls, chunk) -> str:
+        return hex(chunk)[2:].zfill(2)
+
+    @classmethod
+    def change_endian(cls, string) -> str:
+        is_valid, valid_string = cls.validate(string)
+        if not is_valid:
+            raise ValueError("Invalid hex format")
+
+        _output = []
+
+        _chunk_buffer = []
+        for chunk in cls.get_chunked_list(valid_string):
+            _chunk_buffer.append(chunk)
+            if len(_chunk_buffer) == 2:
+                _chunk_buffer.reverse()
+                for _chunk in _chunk_buffer:
+                    _output.append(_chunk)
+                _chunk_buffer = []
+
+        for _chunk in _chunk_buffer:
+            _output.append(_chunk)
+
+        return " ".join(_output)
+
+    @classmethod
+    def show_parsed_results(cls, parsed_results: ParsedResults, depth=0, print_func=print):
+        if parsed_results.has_results:
+            for result in parsed_results.results:
+                if isinstance(result.data, ParsedResults):
+                    print_func("\t" * depth, f"[{result.field}: {result.wire_type}] =>")
+                    cls.show_parsed_results(result.data, depth + 1)
+                else:
+                    print_func("\t" * depth, f"[{result.field}: {result.wire_type}] => {result.data}")
+        if parsed_results.has_remain_data:
+            print_func("\t" * depth, f"left over bytes: {parsed_results.remain_data}")
+
 
 class BytesBuffer:
     def __init__(self):
@@ -254,17 +301,57 @@ def fetch_32bits(self):
         self.set_data_length(4 + 1)
 
 
+class RemainChunkTransaction:
+    def __init__(self):
+        self._is_done = True
+        self._remain_hex_string_list = []
+
+    def consume_chunk(self, chunk):
+        self._remain_hex_string_list.append(
+            Utils.chunk_to_hex_string(chunk)
+        )
+
+    def flush_chunk(self):
+        self._remain_hex_string_list = []
+
+    def start(self):
+        self._is_done = False
+
+    def done(self):
+        self._is_done = True
+        self.flush_chunk()
+
+    @property
+    def is_done(self):
+        return self._is_done
+
+    @property
+    def remain_hex_string_list(self):
+        return self._remain_hex_string_list
+
+    @property
+    def remain_hex_string(self):
+        return " ".join(self._remain_hex_string_list)
+
+    @property
+    def has_remain_data(self):
+        return len(self._remain_hex_string_list) > 0
+
+
 class Parser:
-    def __init__(self, nexted_depth: int = 0):
+    def __init__(self, nexted_depth: int = 0, strict: bool = False):
         self._nested_depth = nexted_depth
         self._buffer = BytesBuffer()
         self._fetcher = Fetcher()
         self._target_field = None
         self._parsed_data: List[ParsedResult] = []
         self._state = State.FIND_FIELD
+        self._is_strict = strict
+
+        self._t = RemainChunkTransaction()
 
     def _create_nested_parser(self) -> Parser:
-        return Parser(nexted_depth=self._nested_depth + 1)
+        return Parser(nexted_depth=self._nested_depth + 1, strict=self._is_strict)
 
     @staticmethod
     def _has_next(chunk_bytes) -> bool:
@@ -294,6 +381,8 @@ def _handler_find_field(self, chunk):
         if self._has_next(chunk):
             return self._next_buffer_handler(value)
 
+        self._t.start()
+
         self._buffer.append(value)
         bit_value = self._get_buffered_value()
         wire_type, field = self._parse_wire_type(bit_value)
@@ -314,7 +403,10 @@ def _handler_find_field(self, chunk):
         elif wire_type == WireType.EGROUP.value:
             self._state = State.TERMINATED
         else:
+            if self._is_strict:
+                raise AssertionError(f"Invalid wire_type: {wire_type}")
             self._state = State.TERMINATED
+
         self._buffer.flush()
 
     def _parse_varint_handler(self, chunk):
@@ -334,6 +426,7 @@ def _parse_varint_handler(self, chunk):
 
         self._state = State.FIND_FIELD
         self._buffer.flush()
+        self._t.done()
 
     def _parse_fixed_handler(self, chunk):
         self._next_buffer_handler(chunk)
@@ -354,6 +447,7 @@ def _parse_fixed_handler(self, chunk):
             self._state = State.FIND_FIELD
             self._buffer.flush()
             self._fetcher.seek()
+            self._t.done()
 
     def _zero_length_delimited_handler(self):
         self._parsed_data.append(
@@ -365,6 +459,7 @@ def _zero_length_delimited_handler(self):
         )
         self._state = State.FIND_FIELD
         self._buffer.flush()
+        self._t.done()
 
     def _parse_length_delimited_handler(self, chunk):
         value = self._get_value(chunk)
@@ -379,6 +474,7 @@ def _parse_length_delimited_handler(self, chunk):
         self._fetcher.set_data_length(data_length)
         self._state = State.GET_DELIMITED_DATA
         self._buffer.flush()
+        self._t.done()
 
     def _next_get_delimited_data_handler(self, value):
         self._fetcher.fetch()
@@ -439,9 +535,16 @@ def _get_delimited_data_handler(self, chunk):
         self._buffer.flush()
         self._fetcher.seek()
         self._state = State.FIND_FIELD
+        self._t.done()
 
     def _create_parsed_results(self) -> ParsedResults:
-        return ParsedResults(results=self._parsed_data)
+        if not self._t.has_remain_data:
+            return ParsedResults(results=self._parsed_data)
+
+        return ParsedResults(
+            results=self._parsed_data,
+            remain_data=self._t.remain_hex_string
+        )
 
     def parse(self, test_target) -> ParsedResults:
         if test_target == "":
@@ -454,6 +557,8 @@ def parse(self, test_target) -> ParsedResults:
         for hex_chunk in Utils.get_chunked_list(validate_string):
             chunk = Utils.hex_string_to_decimal(hex_chunk)
 
+            self._t.consume_chunk(chunk)
+
             if self._state == State.FIND_FIELD:
                 self._handler_find_field(chunk)
 
@@ -478,8 +583,12 @@ def parse(self, test_target) -> ParsedResults:
                 continue
 
             elif self._state == State.TERMINATED:
-                return self._create_parsed_results()
+                pass
+
             else:
                 raise ValueError(f"Unsupported State {self._state}")
 
+        if self._is_strict:
+            assert self._t.is_done, "parsing process is not done, Maybe invalid protobuf"
+
         return self._create_parsed_results()

protobuf_decoder/protobuf_decoder.py

@@ -3,16 +3,6 @@
 from protobuf_decoder.protobuf_decoder import Utils, Parser, ParsedResult, ParsedResults, FixedBitsValue
 
 
-def show_parsed_results(parsed_results: ParsedResults, depth=0):
-    if parsed_results.has_results:
-        for result in parsed_results.results:
-            if isinstance(result.data, ParsedResults):
-                print("\t" * depth, f"[{result.field}: {result.wire_type}] =>")
-                show_parsed_results(result.data, depth + 1)
-            else:
-                print("\t" * depth, f"[{result.field}: {result.wire_type}] => {result.data}")
-
-
 def test_binary_validate():
     is_valid, _ = Utils.validate("08 12")
     assert is_valid is True
@@ -27,6 +17,16 @@ def test_binary_validate():
     assert is_valid is False
 
 
+def test_change_endian():
+    try:
+        Utils.change_endian("081H")
+    except ValueError:
+        assert True
+    assert Utils.change_endian("0812") == "12 08"
+    assert Utils.change_endian("08 12 12") == "12 08 12"
+    assert Utils.change_endian("08 1212 32") == "12 08 32 12"
+
+
 def test_get_chunked_list():
     _, validate_string = Utils.validate("08 12")
     chunked_list = Utils.get_chunked_list(validate_string)
@@ -191,8 +191,12 @@ def test_parser6():
     """
     test_target = " ".join(['ed', '85', '8c', 'ec', '8a', 'a4', 'ed', '8a', 'b8'])
     parsed_data = Parser().parse(test_target)
+
     assert parsed_data.has_results is False
-    assert parsed_data.to_dict() == {'results': []}
+
+    assert parsed_data.to_dict() == {'results': [], 'remain_data': 'ed 85 8c ec 8a a4 ed 8a b8', }
+    assert parsed_data.has_remain_data is True
+    assert parsed_data.remain_data == "ed 85 8c ec 8a a4 ed 8a b8"
 
 
 def test_parser7():
@@ -234,6 +238,7 @@ def test_parser8():
 
     # binary
     0A 04 74 65 73 74 0A 05 74 65 73 74 32
+    0A 04 74 65 73 74 0A 05 74 65 73 74 32 12 01 61 12 01 62
 
     """
     test_target = "0A 04 74 65 73 74 0A 05 74 65 73 74 32"
@@ -1905,3 +1910,61 @@ def test_inner_protobuf_8():
                                                                                ParsedResult(field=2, wire_type='varint',
                                                                                             data=1)])),
                                                               ParsedResult(field=12, wire_type='varint', data=0)]))])
+
+
+def test_strict_protobuf_1():
+    test_target = "02 04 74 65 73 74 02 05 74 65 73 74 32 00 00 00 00 0d 1d"
+    try:
+        Parser(strict=True).parse(test_target)
+    except AssertionError as error:
+        assert "parsing process is not done" in str(error)
+    else:
+        assert False
+
+
+def test_strict_protobuf_2():
+    test_target = "800000000f677270632d7374617475733a300d"
+    try:
+        Parser(strict=True).parse(test_target)
+    except AssertionError as error:
+        assert "Invalid wire_type: 7" in str(error)
+    else:
+        assert False
+
+
+def test_remain_protobuf_1():
+    test_target = "02 04 74 65 73 74 02 05 74 65 73 74 32 00 00 00 00 0d 1d"
+    parsed_data = Parser().parse(test_target)
+    assert parsed_data.has_remain_data
+    assert parsed_data.remain_data == '0d 1d'
+    assert parsed_data == ParsedResults(
+        results=[
+            ParsedResult(field=0, wire_type='string', data='test'),
+            ParsedResult(field=0, wire_type='string', data='test2'),
+            ParsedResult(field=0, wire_type='varint', data=0),
+            ParsedResult(field=0, wire_type='varint', data=0)
+        ], remain_data='0d 1d'
+    )
+    assert parsed_data.to_dict() == {'remain_data': '0d 1d',
+                                     'results': [{'data': 'test', 'field': 0, 'wire_type': 'string'},
+                                                 {'data': 'test2', 'field': 0, 'wire_type': 'string'},
+                                                 {'data': 0, 'field': 0, 'wire_type': 'varint'},
+                                                 {'data': 0, 'field': 0, 'wire_type': 'varint'}]}
+
+
+def test_remain_protobuf_2():
+    test_target = "800000000f677270632d7374617475733a300d"
+    parsed_data = Parser(strict=False).parse(test_target)
+    assert parsed_data.has_remain_data
+    assert parsed_data.remain_data == '67 72 70 63 2d 73 74 61 74 75 73 3a 30 0d'
+
+    assert parsed_data == ParsedResults(
+        results=[
+            ParsedResult(field=0, wire_type='varint', data=0),
+            ParsedResult(field=0, wire_type='varint', data=15)
+        ],
+        remain_data='67 72 70 63 2d 73 74 61 74 75 73 3a 30 0d'
+    )
+    assert parsed_data.to_dict() == {'remain_data': '67 72 70 63 2d 73 74 61 74 75 73 3a 30 0d',
+                                     'results': [{'data': 0, 'field': 0, 'wire_type': 'varint'},
+                                                 {'data': 15, 'field': 0, 'wire_type': 'varint'}]}