354 lines
10 KiB
Python
354 lines
10 KiB
Python
"""
|
||
Фрагментация и сборка изображений SDR Rover Link.
|
||
|
||
Изображение рассматривается как обычная последовательность байтов.
|
||
Модулю неважно, используется JPEG, PNG или другой формат.
|
||
"""
|
||
|
||
from dataclasses import dataclass
|
||
import struct
|
||
from typing import Iterable
|
||
|
||
from protocol.packet import MAX_PAYLOAD_SIZE
|
||
|
||
|
||
# ============================================================
|
||
# Формат заголовка фрагмента изображения
|
||
# ============================================================
|
||
|
||
# I — идентификатор изображения, 4 байта
|
||
# H — номер фрагмента, 2 байта
|
||
# H — общее количество фрагментов, 2 байта
|
||
# I — размер исходного изображения, 4 байта
|
||
FRAGMENT_HEADER_FORMAT = ">IHHI"
|
||
|
||
FRAGMENT_HEADER_SIZE = struct.calcsize(
|
||
FRAGMENT_HEADER_FORMAT
|
||
)
|
||
|
||
MAX_FRAGMENT_DATA_SIZE = (
|
||
MAX_PAYLOAD_SIZE - FRAGMENT_HEADER_SIZE
|
||
)
|
||
|
||
|
||
# ============================================================
|
||
# Ошибки
|
||
# ============================================================
|
||
|
||
class ImageFragmentError(ValueError):
|
||
"""Ошибка структуры фрагмента изображения."""
|
||
|
||
|
||
class MissingFragmentsError(ImageFragmentError):
|
||
"""Отсутствуют необходимые фрагменты изображения."""
|
||
|
||
|
||
# ============================================================
|
||
# Представление одного фрагмента
|
||
# ============================================================
|
||
|
||
@dataclass(frozen=True)
|
||
class ImageFragment:
|
||
"""
|
||
Один фрагмент изображения.
|
||
|
||
Атрибуты
|
||
--------
|
||
image_id:
|
||
Идентификатор изображения.
|
||
|
||
fragment_index:
|
||
Номер фрагмента. Нумерация начинается с нуля.
|
||
|
||
fragment_count:
|
||
Общее количество фрагментов.
|
||
|
||
original_size:
|
||
Размер исходного изображения в байтах.
|
||
|
||
data:
|
||
Часть исходного изображения.
|
||
"""
|
||
|
||
image_id: int
|
||
fragment_index: int
|
||
fragment_count: int
|
||
original_size: int
|
||
data: bytes
|
||
|
||
|
||
# ============================================================
|
||
# Разбиение изображения
|
||
# ============================================================
|
||
|
||
def split_image_bytes(
|
||
image_bytes: bytes,
|
||
image_id: int,
|
||
fragment_data_size: int = 512,
|
||
) -> list[ImageFragment]:
|
||
"""
|
||
Разбить изображение на фрагменты.
|
||
|
||
Параметры
|
||
--------
|
||
image_bytes:
|
||
Полное содержимое файла изображения.
|
||
|
||
image_id:
|
||
Идентификатор изображения от 0 до 4294967295.
|
||
|
||
fragment_data_size:
|
||
Максимальный размер данных одного фрагмента.
|
||
|
||
Возвращает
|
||
----------
|
||
Список объектов ImageFragment.
|
||
"""
|
||
|
||
if not isinstance(image_bytes, (bytes, bytearray)):
|
||
raise TypeError(
|
||
"image_bytes должен иметь тип bytes или bytearray"
|
||
)
|
||
|
||
image_bytes = bytes(image_bytes)
|
||
|
||
if not image_bytes:
|
||
raise ValueError(
|
||
"Изображение не должно быть пустым"
|
||
)
|
||
|
||
if not 0 <= image_id <= 0xFFFFFFFF:
|
||
raise ValueError(
|
||
"image_id должен находиться "
|
||
"в диапазоне 0...4294967295"
|
||
)
|
||
|
||
if not 1 <= fragment_data_size <= MAX_FRAGMENT_DATA_SIZE:
|
||
raise ValueError(
|
||
"Некорректный размер фрагмента"
|
||
)
|
||
|
||
fragment_count = (
|
||
len(image_bytes)
|
||
+ fragment_data_size
|
||
- 1
|
||
) // fragment_data_size
|
||
|
||
if fragment_count > 65535:
|
||
raise ValueError(
|
||
"Изображение требует более 65535 фрагментов"
|
||
)
|
||
|
||
fragments = []
|
||
|
||
for fragment_index in range(fragment_count):
|
||
|
||
start = fragment_index * fragment_data_size
|
||
end = start + fragment_data_size
|
||
|
||
fragment_data = image_bytes[start:end]
|
||
|
||
fragments.append(
|
||
ImageFragment(
|
||
image_id=image_id,
|
||
fragment_index=fragment_index,
|
||
fragment_count=fragment_count,
|
||
original_size=len(image_bytes),
|
||
data=fragment_data,
|
||
)
|
||
)
|
||
|
||
return fragments
|
||
|
||
|
||
# ============================================================
|
||
# Преобразование фрагмента в байты
|
||
# ============================================================
|
||
|
||
def encode_image_fragment(
|
||
fragment: ImageFragment,
|
||
) -> bytes:
|
||
"""
|
||
Преобразовать ImageFragment в полезную нагрузку пакета.
|
||
"""
|
||
|
||
if not isinstance(fragment, ImageFragment):
|
||
raise TypeError(
|
||
"fragment должен иметь тип ImageFragment"
|
||
)
|
||
|
||
if not 0 <= fragment.fragment_index < fragment.fragment_count:
|
||
raise ImageFragmentError(
|
||
"Некорректный номер фрагмента"
|
||
)
|
||
|
||
if not fragment.data:
|
||
raise ImageFragmentError(
|
||
"Фрагмент не должен быть пустым"
|
||
)
|
||
|
||
if len(fragment.data) > MAX_FRAGMENT_DATA_SIZE:
|
||
raise ImageFragmentError(
|
||
"Данные фрагмента слишком велики"
|
||
)
|
||
|
||
header = struct.pack(
|
||
FRAGMENT_HEADER_FORMAT,
|
||
fragment.image_id,
|
||
fragment.fragment_index,
|
||
fragment.fragment_count,
|
||
fragment.original_size,
|
||
)
|
||
|
||
return header + fragment.data
|
||
|
||
|
||
# ============================================================
|
||
# Преобразование байтов обратно во фрагмент
|
||
# ============================================================
|
||
|
||
def decode_image_fragment(
|
||
payload: bytes,
|
||
) -> ImageFragment:
|
||
"""
|
||
Разобрать полезную нагрузку пакета изображения.
|
||
"""
|
||
|
||
if not isinstance(payload, (bytes, bytearray)):
|
||
raise TypeError(
|
||
"payload должен иметь тип bytes или bytearray"
|
||
)
|
||
|
||
payload = bytes(payload)
|
||
|
||
if len(payload) <= FRAGMENT_HEADER_SIZE:
|
||
raise ImageFragmentError(
|
||
"Фрагмент слишком короткий"
|
||
)
|
||
|
||
header = payload[:FRAGMENT_HEADER_SIZE]
|
||
fragment_data = payload[FRAGMENT_HEADER_SIZE:]
|
||
|
||
(
|
||
image_id,
|
||
fragment_index,
|
||
fragment_count,
|
||
original_size,
|
||
) = struct.unpack(
|
||
FRAGMENT_HEADER_FORMAT,
|
||
header,
|
||
)
|
||
|
||
if fragment_count == 0:
|
||
raise ImageFragmentError(
|
||
"Количество фрагментов не может быть равно нулю"
|
||
)
|
||
|
||
if fragment_index >= fragment_count:
|
||
raise ImageFragmentError(
|
||
"Номер фрагмента выходит за допустимый диапазон"
|
||
)
|
||
|
||
if original_size == 0:
|
||
raise ImageFragmentError(
|
||
"Размер исходного изображения равен нулю"
|
||
)
|
||
|
||
return ImageFragment(
|
||
image_id=image_id,
|
||
fragment_index=fragment_index,
|
||
fragment_count=fragment_count,
|
||
original_size=original_size,
|
||
data=fragment_data,
|
||
)
|
||
|
||
|
||
# ============================================================
|
||
# Сборка изображения
|
||
# ============================================================
|
||
|
||
def reassemble_image(
|
||
fragments: Iterable[ImageFragment],
|
||
) -> bytes:
|
||
"""
|
||
Собрать исходное изображение из фрагментов.
|
||
|
||
Фрагменты могут поступать в произвольном порядке.
|
||
"""
|
||
|
||
fragments = list(fragments)
|
||
|
||
if not fragments:
|
||
raise MissingFragmentsError(
|
||
"Список фрагментов пуст"
|
||
)
|
||
|
||
first_fragment = fragments[0]
|
||
|
||
image_id = first_fragment.image_id
|
||
fragment_count = first_fragment.fragment_count
|
||
original_size = first_fragment.original_size
|
||
|
||
fragments_by_index = {}
|
||
|
||
for fragment in fragments:
|
||
|
||
if fragment.image_id != image_id:
|
||
raise ImageFragmentError(
|
||
"Фрагменты относятся к разным изображениям"
|
||
)
|
||
|
||
if fragment.fragment_count != fragment_count:
|
||
raise ImageFragmentError(
|
||
"Различается общее количество фрагментов"
|
||
)
|
||
|
||
if fragment.original_size != original_size:
|
||
raise ImageFragmentError(
|
||
"Различается размер исходного изображения"
|
||
)
|
||
|
||
existing_fragment = fragments_by_index.get(
|
||
fragment.fragment_index
|
||
)
|
||
|
||
if existing_fragment is not None:
|
||
|
||
if existing_fragment.data != fragment.data:
|
||
raise ImageFragmentError(
|
||
"Получены разные данные "
|
||
"для одного номера фрагмента"
|
||
)
|
||
|
||
# Полностью одинаковый повторный фрагмент
|
||
# можно безопасно проигнорировать.
|
||
continue
|
||
|
||
fragments_by_index[
|
||
fragment.fragment_index
|
||
] = fragment
|
||
|
||
missing_indexes = [
|
||
index
|
||
for index in range(fragment_count)
|
||
if index not in fragments_by_index
|
||
]
|
||
|
||
if missing_indexes:
|
||
raise MissingFragmentsError(
|
||
"Отсутствуют фрагменты: "
|
||
+ ", ".join(map(str, missing_indexes))
|
||
)
|
||
|
||
assembled_bytes = b"".join(
|
||
fragments_by_index[index].data
|
||
for index in range(fragment_count)
|
||
)
|
||
|
||
if len(assembled_bytes) != original_size:
|
||
raise ImageFragmentError(
|
||
"Размер собранного изображения "
|
||
"не совпадает с исходным размером"
|
||
)
|
||
|
||
return assembled_bytes |