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):
+class HazelPacketType(bytes, enum.Enum):
-    UNRELIABLE = (0,)
+    UNRELIABLE = bytes([0])
-    RELIABLE = (1,)
+    RELIABLE = bytes([1])
 
-    HELLO = (8,)
+    HELLO = bytes([8])
-    PING = (12,)
+    PING = bytes([12])
-    ACK = (10,)
+    ACK = bytes([10])
-    FIN = (9,)
+    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,8 +82,8 @@ class Parser:
 
         return result
 
-    @staticmethod
+
-    def match_spec(cls: Type[Packet], fields: List[FieldSpec], data: bytes) -> Packet:
+def _match_spec(cls: Type[Packet], fields: List[FieldSpec], data: bytes) -> Packet:
 
     buffer = Buffer(data)
 
@@ -96,8 +92,8 @@ class Parser:
     res_data: List[Any] = list()
 
     for fieldspec in fields:
-            if isinstance(fieldspec, int):
+        if isinstance(fieldspec, bytes):
-                assert buffer.consume(1)[0] == fieldspec
+            assert buffer.consume(len(fieldspec)) == fieldspec
             res_data.append(fieldspec)
             continue
 
@@ -125,3 +121,7 @@ class Parser:
             res_data.append(fielddata)
 
     return cls(res_data, **backref)
+
+
+def unknown(n: int, format: Extractor = None) -> FieldSpec:
+    return (None, n, format)