SDR-Rover/tests/lab007_image_fragmentation.py

"""
Lab007. Разбиение JPEG на фрагменты и обратная сборка.

Этапы:
1. Создание тестового JPEG.
2. Чтение файла как последовательности байтов.
3. Разбиение изображения на фрагменты.
4. Формирование цифрового пакета для каждого фрагмента.
5. Перемешивание порядка доставки пакетов.
6. Проверка CRC и разбор пакетов.
7. Обратная сборка изображения.
8. Сравнение контрольных сумм файлов.
"""

from hashlib import sha256
from pathlib import Path
import random

from PIL import Image

from protocol.image_fragments import (
    decode_image_fragment,
    encode_image_fragment,
    reassemble_image,
    split_image_bytes,
)

from protocol.packet import (
    MESSAGE_TYPE_IMAGE_FRAGMENT,
    build_packet,
    parse_packet,
)


# ============================================================
# Настройки
# ============================================================

IMAGE_WIDTH = 320
IMAGE_HEIGHT = 240

IMAGE_ID = 20260713

FRAGMENT_DATA_SIZE = 512

random_generator = random.Random(2026)


# ============================================================
# Пути к файлам
# ============================================================

source_directory = Path("data/raw")
received_directory = Path("data/received")

source_directory.mkdir(
    parents=True,
    exist_ok=True,
)

received_directory.mkdir(
    parents=True,
    exist_ok=True,
)

source_path = (
    source_directory
    / "lab007_source.jpg"
)

received_path = (
    received_directory
    / "lab007_reassembled.jpg"
)


# ============================================================
# Создание тестового JPEG
# ============================================================

image = Image.new(
    "RGB",
    (IMAGE_WIDTH, IMAGE_HEIGHT),
)

pixels = image.load()

for y in range(IMAGE_HEIGHT):
    for x in range(IMAGE_WIDTH):

        red = (
            x * 255
            // (IMAGE_WIDTH - 1)
        )

        green = (
            y * 255
            // (IMAGE_HEIGHT - 1)
        )

        blue = (
            (x + y) * 255
            // (
                IMAGE_WIDTH
                + IMAGE_HEIGHT
                - 2
            )
        )

        pixels[x, y] = (
            red,
            green,
            blue,
        )

image.save(
    source_path,
    format="JPEG",
    quality=60,
    optimize=True,
)


# ============================================================
# Чтение изображения
# ============================================================

source_bytes = source_path.read_bytes()

source_hash = sha256(
    source_bytes
).hexdigest()


# ============================================================
# Разбиение изображения на фрагменты
# ============================================================

fragments = split_image_bytes(
    image_bytes=source_bytes,
    image_id=IMAGE_ID,
    fragment_data_size=FRAGMENT_DATA_SIZE,
)


# ============================================================
# Формирование цифровых пакетов
# ============================================================

packets = []

for fragment in fragments:

    fragment_payload = encode_image_fragment(
        fragment
    )

    packet = build_packet(
        payload=fragment_payload,
        message_type=MESSAGE_TYPE_IMAGE_FRAGMENT,
        sequence_number=fragment.fragment_index,
    )

    packets.append(packet)


# ============================================================
# Имитация доставки в произвольном порядке
# ============================================================

random_generator.shuffle(packets)

received_fragments = []

for packet in packets:

    parsed_packet = parse_packet(packet)

    assert (
        parsed_packet.message_type
        == MESSAGE_TYPE_IMAGE_FRAGMENT
    )

    fragment = decode_image_fragment(
        parsed_packet.payload
    )

    assert (
        parsed_packet.sequence_number
        == fragment.fragment_index
    )

    received_fragments.append(fragment)


# ============================================================
# Сборка изображения
# ============================================================

received_bytes = reassemble_image(
    received_fragments
)

received_path.write_bytes(
    received_bytes
)

received_hash = sha256(
    received_bytes
).hexdigest()


# ============================================================
# Проверки
# ============================================================

assert received_bytes == source_bytes
assert received_hash == source_hash

with Image.open(received_path) as received_image:
    received_image.verify()


# ============================================================
# Результаты
# ============================================================

print("=== Lab007. Фрагментация изображения ===")

print("\nИсходный файл:")
print(source_path)

print("\nВосстановленный файл:")
print(received_path)

print("\nРазмер изображения:")
print(
    IMAGE_WIDTH,
    "x",
    IMAGE_HEIGHT,
    "пикселей",
)

print("\nРазмер JPEG:")
print(
    len(source_bytes),
    "байт",
)

print("\nРазмер данных одного фрагмента:")
print(
    FRAGMENT_DATA_SIZE,
    "байт",
)

print("\nКоличество фрагментов:")
print(
    len(fragments)
)

print("\nSHA-256 исходного файла:")
print(source_hash)

print("\nSHA-256 восстановленного файла:")
print(received_hash)

print(
    "\nПроверка пройдена: "
    "JPEG восстановлен байт в байт."
)