diff --git a/growtopia/_types.py b/growtopia/_types.py index f5bd74f..a428c27 100644 --- a/growtopia/_types.py +++ b/growtopia/_types.py @@ -7,17 +7,21 @@ "int8", "int16", "int32", + "TVariantValue", "TVariant", ) from typing import ( + TYPE_CHECKING, Generic, TypeVar, - TYPE_CHECKING ) if TYPE_CHECKING: - from growtopia.net.protocol.variant import Variant + from growtopia.net.protocol.variant import ( + IntVariant, + StrVariant, + ) type LengthPrefixedStr = str type LengthPrefixedData = bytes | bytearray @@ -26,9 +30,12 @@ type int8 = int type int16 = int type int32 = int -type TVariant = int | float | str | tuple | Variant +type TVariantValue = int | str +type TVariant = IntVariant | StrVariant T = TypeVar("T") + + class Pack(Generic[T]): pass diff --git a/growtopia/net/protocol/__init__.py b/growtopia/net/protocol/__init__.py index 4c9be19..f345b48 100644 --- a/growtopia/net/protocol/__init__.py +++ b/growtopia/net/protocol/__init__.py @@ -1,2 +1,3 @@ from .enums import * from .packet import * +from .variant import * diff --git a/growtopia/net/protocol/enums.py b/growtopia/net/protocol/enums.py index abc3ada..dc5cced 100644 --- a/growtopia/net/protocol/enums.py +++ b/growtopia/net/protocol/enums.py @@ -2,6 +2,7 @@ "PacketType", "UpdateType", "UpdateFlags", + "VariantType", ) from aenum import ( @@ -64,3 +65,12 @@ class UpdateType(IntEnum): class UpdateFlags(IntFlag): NONE = 0 EXTRA_DATA = 1 << 3 + + +class VariantType(IntEnum): + FLOAT = 1 + STR = 2 + VECTOR2 = 3 + VECTOR3 = 4 + UINT = 8 + INT = 9 diff --git a/growtopia/net/protocol/packet.py b/growtopia/net/protocol/packet.py index 2112317..6fedd8d 100644 --- a/growtopia/net/protocol/packet.py +++ b/growtopia/net/protocol/packet.py @@ -15,17 +15,24 @@ AllStr, OptionalPack, Pack, + TVariant, int8, int32, ) +from growtopia.net.protocol.variant import ( + TYPE_TO_OBJ_MAPPING, +) from growtopia.utils import ( + LOG_LEVEL_ERROR, Packer, + log, ) from .enums import ( PacketType, UpdateFlags, UpdateType, + VariantType, ) @@ -176,3 +183,42 @@ def enet_packet(self, flags: int = enet.PACKET_FLAG_RELIABLE) -> enet.Packet: return enet.Packet(self._prepacked_data, flags) return enet.Packet(self.pack(), flags) + + def get_variant_list(self) -> list[TVariant]: + if not self.flags & UpdateFlags.EXTRA_DATA: + return [] + + offset = 1 + variant_count = self.extra_data[0] + variants = [] + + for i in range(variant_count): + variant_index = self.extra_data[offset] + variant_type = VariantType(self.extra_data[offset + 1]) + + if variant_index != i: + log( + LOG_LEVEL_ERROR, + f"Variant ({variant_type.name}) index doesn't match its position in the list!! (expected index: {i}, got {variant_index})", + ) + + return [] + + variant_obj = TYPE_TO_OBJ_MAPPING[variant_type]().from_bytes(self.extra_data[offset:]) + variants.append(variant_obj) + + offset += variant_obj.get_size() + + return variants + + def set_variant_list(self, *variants: TVariant, keep_index: bool = False) -> None: + self.flags |= UpdateFlags.EXTRA_DATA + self.extra_data = bytearray(len(variants).to_bytes(1, "little")) + + for i, variant in enumerate(variants): + if not keep_index: + variant.index = i + + self.extra_data.extend(variant.pack()) + + self.extra_data_size = len(self.extra_data) diff --git a/growtopia/net/protocol/variant.py b/growtopia/net/protocol/variant.py new file mode 100644 index 0000000..f3e5d12 --- /dev/null +++ b/growtopia/net/protocol/variant.py @@ -0,0 +1,60 @@ +__all__ = ( + "Variant", + "IntVariant", + "StrVariant", + "TYPE_TO_OBJ_MAPPING", +) + +from growtopia._types import ( + LengthPrefixedStr, + Pack, + TVariant, + TVariantValue, + int8, + int32, +) +from growtopia.net.protocol.enums import ( + VariantType, +) +from growtopia.utils import ( + Packer, +) + + +class Variant: + def __init__( + self, index: int, variant_type: VariantType, variant_value: TVariantValue | None = None + ) -> None: + self.index: int = index + self.variant_type: VariantType = variant_type + self.variant_value: TVariantValue | None = variant_value + + def get_size(self) -> int: + if self.variant_type == VariantType.STR: + return 2 + len(self.variant_value) + if self.variant_type == VariantType.INT: + return 6 + + +class IntVariant(Variant, Packer): + index: Pack[int8] + variant_type: Pack[int8] + variant_value: Pack[int32] + + def __init__(self, variant_value: int | None = None) -> None: + Variant.__init__(self, -1, VariantType.INT, variant_value) + + +class StrVariant(Variant, Packer): + index: Pack[int8] + variant_type: Pack[int8] + variant_value: Pack[LengthPrefixedStr] + + def __init__(self, variant_value: str | None = None) -> None: + Variant.__init__(self, -1, VariantType.STR, variant_value) + + +TYPE_TO_OBJ_MAPPING: dict[VariantType, TVariant] = { + VariantType.INT: IntVariant, + VariantType.STR: StrVariant, +} diff --git a/growtopia/utils/packer.py b/growtopia/utils/packer.py index edb7d40..8b4f7b0 100644 --- a/growtopia/utils/packer.py +++ b/growtopia/utils/packer.py @@ -2,9 +2,9 @@ from typing import ( Callable, + Self, Type, get_origin, - Self, ) from growtopia._types import (