ForcePilot/backend/package/yuxi/channel/extensions/minecraft/protocol.py

import struct

MAX_VARINT_BYTES = 5
MAX_VARLONG_BYTES = 10
SEGMENT_BITS = 0x7F
CONTINUE_BIT = 0x80


def read_varint(stream: bytes, offset: int = 0) -> tuple[int, int]:
    value = 0
    position = 0
    consumed = 0
    while True:
        if offset + consumed >= len(stream):
            raise ValueError("VarInt: unexpected end of stream")
        byte = stream[offset + consumed]
        value |= (byte & SEGMENT_BITS) << position
        consumed += 1
        if not (byte & CONTINUE_BIT):
            break
        position += 7
        if position >= 32:
            raise ValueError("VarInt too large (> 5 bytes)")
    return value, consumed


def write_varint(value: int) -> bytes:
    result = bytearray()
    value = value & 0xFFFFFFFF
    while True:
        if value & ~SEGMENT_BITS:
            result.append((value & SEGMENT_BITS) | CONTINUE_BIT)
            value >>= 7
        else:
            result.append(value & SEGMENT_BITS)
            break
    return bytes(result)


def read_varlong(stream: bytes, offset: int = 0) -> tuple[int, int]:
    value = 0
    position = 0
    consumed = 0
    while True:
        if offset + consumed >= len(stream):
            raise ValueError("VarLong: unexpected end of stream")
        byte = stream[offset + consumed]
        value |= (byte & SEGMENT_BITS) << position
        consumed += 1
        if not (byte & CONTINUE_BIT):
            break
        position += 7
        if position >= 64:
            raise ValueError("VarLong too large (> 10 bytes)")
    return value, consumed


def write_varlong(value: int) -> bytes:
    result = bytearray()
    value = value & 0xFFFFFFFFFFFFFFFF
    while True:
        if value & ~SEGMENT_BITS:
            result.append((value & SEGMENT_BITS) | CONTINUE_BIT)
            value >>= 7
        else:
            result.append(value & SEGMENT_BITS)
            break
    return bytes(result)


def read_string(stream: bytes, offset: int) -> tuple[str, int]:
    length, consumed = read_varint(stream, offset)
    if length == 0:
        return "", consumed
    start = offset + consumed
    end = start + length
    if end > len(stream):
        raise ValueError(f"String: expected {length} bytes but only {len(stream) - start} available")
    return stream[start:end].decode("utf-8", errors="replace"), consumed + length


def write_string(value: str) -> bytes:
    encoded = value.encode("utf-8")
    return write_varint(len(encoded)) + encoded


def read_uuid(stream: bytes, offset: int) -> tuple[str, int]:
    if offset + 16 > len(stream):
        raise ValueError("UUID: unexpected end of stream")
    raw = stream[offset : offset + 16]
    uuid_str = "-".join(
        [
            raw[0:4].hex(),
            raw[4:6].hex(),
            raw[6:8].hex(),
            raw[8:10].hex(),
            raw[10:16].hex(),
        ]
    )
    return uuid_str, 16


def write_uuid(uuid_str: str) -> bytes:
    hex_str = uuid_str.replace("-", "")
    return bytes.fromhex(hex_str)


def read_byte(stream: bytes, offset: int) -> tuple[int, int]:
    if offset >= len(stream):
        raise ValueError("Byte: unexpected end of stream")
    return stream[offset], 1


def write_byte(value: int) -> bytes:
    return bytes([value & 0xFF])


def read_short(stream: bytes, offset: int) -> tuple[int, int]:
    return struct.unpack(">h", stream[offset : offset + 2])[0], 2


def write_short(value: int) -> bytes:
    return struct.pack(">h", value)


def read_ushort(stream: bytes, offset: int) -> tuple[int, int]:
    return struct.unpack(">H", stream[offset : offset + 2])[0], 2


def write_ushort(value: int) -> bytes:
    return struct.pack(">H", value)


def read_int(stream: bytes, offset: int) -> tuple[int, int]:
    return struct.unpack(">i", stream[offset : offset + 4])[0], 4


def write_int(value: int) -> bytes:
    return struct.pack(">i", value)


def read_long(stream: bytes, offset: int) -> tuple[int, int]:
    return struct.unpack(">q", stream[offset : offset + 8])[0], 8


def write_long(value: int) -> bytes:
    return struct.pack(">q", value)


def read_double(stream: bytes, offset: int) -> tuple[float, int]:
    return struct.unpack(">d", stream[offset : offset + 8])[0], 8


def write_double(value: float) -> bytes:
    return struct.pack(">d", value)


def read_float(stream: bytes, offset: int) -> tuple[float, int]:
    return struct.unpack(">f", stream[offset : offset + 4])[0], 4


def write_float(value: float) -> bytes:
    return struct.pack(">f", value)


def read_bool(stream: bytes, offset: int) -> tuple[bool, int]:
    val, consumed = read_byte(stream, offset)
    return val != 0, consumed


def write_bool(value: bool) -> bytes:
    return b"\x01" if value else b"\x00"


def read_packet_frame(stream: bytes, offset: int = 0) -> tuple[int, int, int, int]:
    packet_length, len_consumed = read_varint(stream, offset)
    data_start = offset + len_consumed
    packet_id, id_consumed = read_varint(stream, data_start)
    data_offset = data_start + id_consumed
    data_length = packet_length - id_consumed
    total_consumed = len_consumed + packet_length
    return packet_id, data_offset, data_length, total_consumed


def write_packet_frame(packet_id: int, data: bytes = b"") -> bytes:
    packet_id_bytes = write_varint(packet_id)
    payload = packet_id_bytes + data
    length_bytes = write_varint(len(payload))
    return length_bytes + payload


def read_varint_at(stream: bytes, offset: int) -> tuple[int, int]:
    value, consumed = read_varint(stream, offset)
    return value, offset + consumed


def read_string_at(stream: bytes, offset: int) -> tuple[str, int]:
    length, new_offset = read_varint(stream, offset)
    start = new_offset
    end = start + length
    if end > len(stream):
        raise ValueError(f"String: expected {length} bytes but only {len(stream) - start} available")
    return stream[start:end].decode("utf-8", errors="replace"), end