Refactor

proto.py

@@ -2,19 +2,19 @@ import enum
 import binascii
 
 import protoparser
-from protoparser import Packet
+from protoparser import Packet, unknown
 
 
-class HazelPacketType(enum.IntEnum):
-    UNRELIABLE = (0,)
-    RELIABLE = (1,)
+class HazelPacketType(bytes, enum.Enum):
+    UNRELIABLE = bytes([0])
+    RELIABLE = bytes([1])
 
-    HELLO = (8,)
-    PING = (12,)
-    ACK = (10,)
-    FIN = (9,)
+    HELLO = bytes([8])
+    PING = bytes([12])
+    ACK = bytes([10])
+    FIN = bytes([9])
 
-    FRAGMENT = (11,)  # not observed yet, maybe unused in among us?
+    FRAGMENT = bytes([11])  # not observed yet, maybe unused in among us?
 
 
 def int_big_endian(data: bytes) -> int:
@@ -34,7 +34,7 @@ class PingPacket(Packet):
         return f"Ping {self.nonce}"
 
 
-@Parser.register(HazelPacketType.ACK, ("nonce", 2, int_big_endian), 0xFF)
+@Parser.register(HazelPacketType.ACK, ("nonce", 2, int_big_endian), b"\xFF")
 class AckPacket(Packet):
     def __init__(self, data, nonce):
         self.nonce = nonce
@@ -52,7 +52,7 @@ class FinPacket(Packet):
 
 @Parser.register(
     HazelPacketType.HELLO,
-    (None, 7, None),
+    unknown(7),
     ("name_len", 1, int_big_endian),
     ("name", "name_len", bytes.decode),
 )

protoparser.py

@@ -5,13 +5,9 @@ from typing import Tuple, Dict, List, Union, Callable, Optional, Any, Type
 Extractor = Callable[[bytes], Any]
 
 
-def int_big_endian(data: bytes) -> int:
-    return int.from_bytes(data, "big")
-
-
 FieldSpec = Union[
     # specific value to be expected
-    int,
+    bytes,
     # [named] field with fixed length
     Tuple[Optional[str], int, Optional[Extractor]],
     # [named] field with length backreference
@@ -73,7 +69,7 @@ class Parser:
         result = None
         for (cls, fields) in self.specs:
             try:
-                m = self.match_spec(cls, fields, data)
+                m = _match_spec(cls, fields, data)
             except AssertionError:
                 continue
             if m:
@@ -86,42 +82,46 @@ class Parser:
 
         return result
 
-    @staticmethod
-    def match_spec(cls: Type[Packet], fields: List[FieldSpec], data: bytes) -> Packet:
 
-        buffer = Buffer(data)
+def _match_spec(cls: Type[Packet], fields: List[FieldSpec], data: bytes) -> Packet:
 
-        backref: Dict[str, Any] = dict()
+    buffer = Buffer(data)
 
-        res_data: List[Any] = list()
+    backref: Dict[str, Any] = dict()
 
-        for fieldspec in fields:
-            if isinstance(fieldspec, int):
-                assert buffer.consume(1)[0] == fieldspec
-                res_data.append(fieldspec)
-                continue
+    res_data: List[Any] = list()
 
-            if fieldspec is None:
-                res_data.append(buffer.data)
-                break  # TODO implement unknown blob can also be in the middle
+    for fieldspec in fields:
+        if isinstance(fieldspec, bytes):
+            assert buffer.consume(len(fieldspec)) == fieldspec
+            res_data.append(fieldspec)
+            continue
 
-            if isinstance(fieldspec, tuple):
-                fieldname, fieldlen, extractor = fieldspec
+        if fieldspec is None:
+            res_data.append(buffer.data)
+            break  # TODO implement unknown blob can also be in the middle
 
-                # backreference
-                if isinstance(fieldlen, str):
-                    fieldlen = backref[fieldlen]
+        if isinstance(fieldspec, tuple):
+            fieldname, fieldlen, extractor = fieldspec
 
-                assert isinstance(fieldlen, int)
+            # backreference
+            if isinstance(fieldlen, str):
+                fieldlen = backref[fieldlen]
 
-                fielddata = buffer.consume(fieldlen)
+            assert isinstance(fieldlen, int)
 
-                if extractor:
-                    fielddata = extractor(fielddata)
+            fielddata = buffer.consume(fieldlen)
 
-                if fieldname:
-                    backref[fieldname] = fielddata
+            if extractor:
+                fielddata = extractor(fielddata)
 
-                res_data.append(fielddata)
+            if fieldname:
+                backref[fieldname] = fielddata
 
-        return cls(res_data, **backref)
+            res_data.append(fielddata)
+
+    return cls(res_data, **backref)
+
+
+def unknown(n: int, format: Extractor = None) -> FieldSpec:
+    return (None, n, format)