bta-proxy/bta_proxy/packets/base.py

76 lines
2.7 KiB
Python
Raw Normal View History

2023-08-24 15:29:57 +03:00
from typing import Any, ClassVar, Type
2023-08-24 17:40:23 +03:00
from bta_proxy.entitydata import EntityData
2023-08-25 21:49:10 +03:00
from ..datainputstream import AsyncDataInputStream
2023-08-24 15:29:57 +03:00
class Packet:
REGISTRY: ClassVar[dict[int, Type['Packet']]] = {}
FIELDS: ClassVar[list[tuple[str, Any]]] = []
2023-08-24 17:40:23 +03:00
packet_id: int
2023-08-24 15:29:57 +03:00
def __init__(self, **params):
for k, v in params.items():
setattr(self, k, v)
@classmethod
2023-08-25 21:49:10 +03:00
async def read_data_from(cls, stream: AsyncDataInputStream) -> 'Packet':
2023-08-24 15:29:57 +03:00
fields: dict = {}
for key, datatype in cls.FIELDS:
2023-08-25 21:49:10 +03:00
fields[key] = await cls.read_field(stream, datatype)
2023-08-24 15:29:57 +03:00
return cls(**fields)
@staticmethod
2023-08-25 21:49:10 +03:00
async def read_field(stream: AsyncDataInputStream, datatype: Any):
2023-08-24 15:29:57 +03:00
match datatype:
case 'uint':
2023-08-25 21:49:10 +03:00
return await stream.read_uint()
2023-08-24 15:29:57 +03:00
case 'int':
2023-08-25 21:49:10 +03:00
return await stream.read_int()
2023-08-24 15:29:57 +03:00
case 'str':
2023-08-25 21:49:10 +03:00
return await stream.read_string()
2023-08-24 15:29:57 +03:00
case 'str', length:
2023-08-25 21:49:10 +03:00
return (await stream.read_string())[:length]
2023-08-24 15:29:57 +03:00
case 'ulong':
2023-08-25 21:49:10 +03:00
return await stream.read_ulong()
2023-08-24 15:29:57 +03:00
case 'long':
2023-08-25 21:49:10 +03:00
return await stream.read_long()
2023-08-24 15:29:57 +03:00
case 'ushort':
2023-08-25 21:49:10 +03:00
return await stream.read_ushort()
2023-08-24 15:29:57 +03:00
case 'short':
2023-08-25 21:49:10 +03:00
return await stream.read_short()
2023-08-24 15:29:57 +03:00
case 'byte':
2023-08-25 21:49:10 +03:00
return await stream.read()
2023-08-24 15:29:57 +03:00
case 'float':
2023-08-25 21:49:10 +03:00
return await stream.read_float()
2023-08-24 15:29:57 +03:00
case 'double':
2023-08-25 21:49:10 +03:00
return await stream.read_double()
2023-08-24 15:29:57 +03:00
case 'bool':
2023-08-25 21:49:10 +03:00
return await stream.read_boolean()
2023-08-24 15:29:57 +03:00
case 'bytes', length:
2023-08-25 21:49:10 +03:00
return await stream.read_bytes(length)
2023-08-24 17:40:23 +03:00
case 'entitydata':
2023-08-25 21:49:10 +03:00
return await EntityData.read_from(stream)
2023-08-24 15:29:57 +03:00
case _:
raise ValueError(f'unknown type {datatype}')
2023-08-24 17:40:23 +03:00
def __init_subclass__(cls, packet_id: int, **kwargs) -> None:
Packet.REGISTRY[packet_id] = cls
cls.packet_id = packet_id
super().__init_subclass__(**kwargs)
2023-08-24 15:29:57 +03:00
@classmethod
2023-08-25 21:49:10 +03:00
async def read_packet(cls, stream: AsyncDataInputStream) -> 'Packet':
packet_id: int = await stream.read()
2023-08-24 15:29:57 +03:00
if packet_id not in cls.REGISTRY:
raise ValueError(f'invalid packet 0x{packet_id:02x}')
2023-08-25 21:49:10 +03:00
pkt = await cls.REGISTRY[packet_id].read_data_from(stream)
pkt.packet_id = packet_id
return pkt
2023-08-24 15:29:57 +03:00
def __repr__(self):
2023-08-24 17:40:23 +03:00
pkt_name = self.REGISTRY[self.packet_id].__name__
2023-08-24 15:29:57 +03:00
fields = []
for name, _ in self.FIELDS:
fields.append(f'{name}={getattr(self, name)!r}')
2023-08-25 21:49:10 +03:00
return f'<{pkt_name} {str.join(", ", fields)}>'