SDR-Rover/protocol/packet.py

"""
Пакетный протокол SDR Rover Link.

Модуль содержит функции:
- build_packet() — формирование пакета;
- parse_packet() — проверка и разбор пакета.

Версия протокола: 0.1
"""

from dataclasses import dataclass
import struct
import zlib


# ============================================================
# Константы протокола
# ============================================================

SYNC_WORD = 0xAA55
PROTOCOL_VERSION = 1

MESSAGE_TYPE_TEXT = 1
MESSAGE_TYPE_TELEMETRY = 2
MESSAGE_TYPE_IMAGE_FRAGMENT = 3
MESSAGE_TYPE_ACK = 4

# Структура заголовка:
#
# H — SYNC, 2 байта
# B — версия, 1 байт
# B — тип сообщения, 1 байт
# H — порядковый номер, 2 байта
# H — длина полезной нагрузки, 2 байта
HEADER_FORMAT = ">HBBHH"

# CRC-32 хранится как беззнаковое 32-битное число.
CRC_FORMAT = ">I"

HEADER_SIZE = struct.calcsize(HEADER_FORMAT)
CRC_SIZE = struct.calcsize(CRC_FORMAT)

MAX_PAYLOAD_SIZE = 65535


# ============================================================
# Ошибки протокола
# ============================================================

class PacketError(ValueError):
    """Общая ошибка структуры пакета."""


class CRCError(PacketError):
    """Ошибка контрольной суммы CRC."""


# ============================================================
# Результат разбора пакета
# ============================================================

@dataclass(frozen=True)
class ParsedPacket:
    """
    Представление успешно принятого пакета.

    Атрибуты:
        version: версия протокола;
        message_type: тип сообщения;
        sequence_number: порядковый номер;
        payload: полезная нагрузка в байтах.
    """

    version: int
    message_type: int
    sequence_number: int
    payload: bytes


# ============================================================
# Формирование пакета
# ============================================================

def build_packet(
    payload: bytes,
    message_type: int,
    sequence_number: int,
) -> bytes:
    """
    Сформировать пакет SDR Rover Link.

    Параметры:
        payload:
            Полезная нагрузка в байтах.

        message_type:
            Тип сообщения от 0 до 255.

        sequence_number:
            Порядковый номер от 0 до 65535.

    Возвращает:
        Полный пакет в виде bytes.
    """

    if not isinstance(payload, (bytes, bytearray)):
        raise TypeError(
            "payload должен иметь тип bytes или bytearray"
        )

    payload = bytes(payload)

    if not 0 <= message_type <= 255:
        raise ValueError(
            "message_type должен находиться в диапазоне 0...255"
        )

    if not 0 <= sequence_number <= 65535:
        raise ValueError(
            "sequence_number должен находиться в диапазоне 0...65535"
        )

    payload_length = len(payload)

    if payload_length > MAX_PAYLOAD_SIZE:
        raise ValueError(
            f"Полезная нагрузка превышает {MAX_PAYLOAD_SIZE} байт"
        )

    header = struct.pack(
        HEADER_FORMAT,
        SYNC_WORD,
        PROTOCOL_VERSION,
        message_type,
        sequence_number,
        payload_length,
    )

    packet_without_crc = header + payload

    crc_value = (
        zlib.crc32(packet_without_crc) & 0xFFFFFFFF
    )

    crc_bytes = struct.pack(
        CRC_FORMAT,
        crc_value,
    )

    return packet_without_crc + crc_bytes


# ============================================================
# Разбор пакета
# ============================================================

def parse_packet(packet: bytes) -> ParsedPacket:
    """
    Проверить и разобрать пакет SDR Rover Link.

    Проверяются:
        - минимальный размер;
        - маркер SYNC;
        - версия протокола;
        - длина пакета;
        - CRC-32.

    Возвращает:
        ParsedPacket.

    Вызывает:
        PacketError — при ошибке структуры;
        CRCError — при несовпадении CRC.
    """

    if not isinstance(packet, (bytes, bytearray)):
        raise TypeError(
            "packet должен иметь тип bytes или bytearray"
        )

    packet = bytes(packet)

    minimum_packet_size = HEADER_SIZE + CRC_SIZE

    if len(packet) < minimum_packet_size:
        raise PacketError(
            "Пакет слишком короткий"
        )

    header = packet[:HEADER_SIZE]

    (
        sync_word,
        version,
        message_type,
        sequence_number,
        payload_length,
    ) = struct.unpack(
        HEADER_FORMAT,
        header,
    )

    if sync_word != SYNC_WORD:
        raise PacketError(
            f"Неверный SYNC: 0x{sync_word:04X}"
        )

    if version != PROTOCOL_VERSION:
        raise PacketError(
            f"Неподдерживаемая версия протокола: {version}"
        )

    expected_packet_size = (
        HEADER_SIZE
        + payload_length
        + CRC_SIZE
    )

    if len(packet) != expected_packet_size:
        raise PacketError(
            "Длина пакета не соответствует полю LENGTH: "
            f"ожидалось {expected_packet_size} байт, "
            f"получено {len(packet)} байт"
        )

    data_without_crc = packet[:-CRC_SIZE]
    received_crc_bytes = packet[-CRC_SIZE:]

    (received_crc,) = struct.unpack(
        CRC_FORMAT,
        received_crc_bytes,
    )

    calculated_crc = (
        zlib.crc32(data_without_crc) & 0xFFFFFFFF
    )

    if received_crc != calculated_crc:
        raise CRCError(
            "CRC не совпадает: "
            f"получено 0x{received_crc:08X}, "
            f"вычислено 0x{calculated_crc:08X}"
        )

    payload = packet[
        HEADER_SIZE:-CRC_SIZE
    ]

    return ParsedPacket(
        version=version,
        message_type=message_type,
        sequence_number=sequence_number,
        payload=payload,
    )