Serdtse-Sysoly/finals/kot1/solve/solve.py

#!/usr/bin/env python
"""Эталонный решатель задачи 1 «В глазах всё написано» (формат CTP1).

Независимая реализация: только stdlib + Pillow + numpy. Не импортирует код
генератора и не читает файлы организатора.

Схема (восстанавливается из публичных улик в метаданных PNG):
  * сетка блоков 2x2 от левого верхнего угла, обход по строкам;
  * блок пригоден, если 4 цвета RGB попарно различны и размах каждого
    канала <= 3 (пригодность инвариантна к перестановке пикселей);
  * канонический порядок цветов: сортировка по ключу (77*R+150*G+29*B, R, G, B);
  * 24 перестановки (0,1,2,3) нумеруются лексикографически от 0,
    номера 0..15 кодируют полубайт, старший полубайт байта первым;
  * первые 16 пригодных блоков — калибровка (номера 0..15 по порядку);
  * далее конверт: magic "CTP1", version=1, payload_length u32 LE,
    payload, crc32 u32 LE от payload.

Использование: python solve_task1.py <participant.png>
Печатает флаг (первая строка payload) в stdout. При ошибке — сообщение
в stderr и ненулевой код возврата.
"""

import struct
import sys
import zlib

import numpy as np
from PIL import Image

RANGE_MAX = 3
CALIBRATION_BLOCKS = 16
MAGIC = b"CTP1"
VERSION = 1
MAX_PAYLOAD = 1 << 20  # 1 МиБ

# Все 24 перестановки (0,1,2,3) в лексикографическом порядке.
LEX_PERMS = sorted(
    (a, b, c, d)
    for a in range(4)
    for b in range(4)
    for c in range(4)
    for d in range(4)
    if len({a, b, c, d}) == 4
)
PERM_INDEX = {p: i for i, p in enumerate(LEX_PERMS)}


class SolveError(Exception):
    """Ожидаемая ошибка решения: повреждённые или чужие данные."""


def suitable(px):
    """Пригодность блока 2x2. px — массив (4, 3) RGB в порядке TL,TR,BL,BR."""
    for i in range(4):
        for j in range(i + 1, 4):
            if px[i][0] == px[j][0] and px[i][1] == px[j][1] and px[i][2] == px[j][2]:
                return False
    lo = px.min(axis=0)
    hi = px.max(axis=0)
    return bool(((hi - lo) <= RANGE_MAX).all())


def block_nibble(px):
    """Номер лексикографической перестановки блока (0..23).

    perm[k] = позиция k-го канонического цвета в порядке чтения блока.
    """
    order = sorted(
        range(4), key=lambda i: (77 * px[i][0] + 150 * px[i][1] + 29 * px[i][2],
                                 px[i][0], px[i][1], px[i][2])
    )
    perm = [0] * 4
    for k, i in enumerate(order):  # i — позиция в порядке чтения цвета C[k]
        perm[k] = i
    return PERM_INDEX[tuple(perm)]


def solve(path):
    img = Image.open(path).convert("RGB")
    w, h = img.size
    arr = np.asarray(img, dtype=np.int32)

    # Пригодные блоки 2x2 в порядке обхода по строкам; неполные краевые
    # блоки игнорируются.
    blocks = []  # (bx, by) левых верхних углов
    nibbles = []
    for by in range(0, h - 1, 2):
        for bx in range(0, w - 1, 2):
            px = arr[by:by + 2, bx:bx + 2].reshape(4, 3)
            if suitable(px):
                blocks.append((bx, by))
                nibbles.append(block_nibble(px))

    if len(blocks) < CALIBRATION_BLOCKS:
        raise SolveError(
            f"слишком мало пригодных блоков: {len(blocks)}, нужно минимум 16"
        )

    # Калибровка: первые 16 пригодных блоков обязаны давать 0..15 по порядку.
    for i in range(CALIBRATION_BLOCKS):
        if nibbles[i] != i:
            raise SolveError(
                f"калибровка нарушена: блок {i} даёт номер {nibbles[i]}, ожидался {i}"
            )

    def read_byte(idx):
        """Байт idx конверта: старший полубайт первым."""
        vals = []
        for half in range(2):
            pos = CALIBRATION_BLOCKS + 2 * idx + half
            if pos >= len(nibbles):
                raise SolveError("конверт усечён: не хватает пригодных блоков")
            n = nibbles[pos]
            if n > 15:
                raise SolveError(
                    f"недопустимый номер перестановки {n} (допустимы 0..15) "
                    f"в блоке {pos}"
                )
            vals.append(n)
        return (vals[0] << 4) | vals[1]

    header = bytes(read_byte(i) for i in range(9))
    if header[:4] != MAGIC:
        raise SolveError(f"неверный magic: {header[:4]!r}, ожидался {MAGIC!r}")
    if header[4] != VERSION:
        raise SolveError(f"неподдерживаемая версия конверта: {header[4]}")
    (plen,) = struct.unpack("<I", header[5:9])
    if plen > MAX_PAYLOAD:
        raise SolveError(
            f"payload_length {plen} превышает лимит {MAX_PAYLOAD} байт"
        )
    total = 13 + plen
    if CALIBRATION_BLOCKS + 2 * total > len(blocks):
        raise SolveError(
            f"конверт усечён: нужно {total} байт, "
            f"блоков хватает на {(len(blocks) - CALIBRATION_BLOCKS) // 2}"
        )

    rest = bytes(read_byte(i) for i in range(9, total))
    payload = rest[:plen]
    (crc_want,) = struct.unpack("<I", rest[plen:])
    crc_got = zlib.crc32(payload) & 0xFFFFFFFF
    if crc_got != crc_want:
        raise SolveError(
            f"несовпадение crc32: вычислено {crc_got:08x}, записано {crc_want:08x}"
        )
    return payload


def main(argv):
    if len(argv) != 2:
        print(f"usage: {argv[0]} <participant.png>", file=sys.stderr)
        return 2
    try:
        payload = solve(argv[1])
    except SolveError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    except Exception as exc:  # битый файл, не-PNG и т.п.
        print(f"error: не удалось обработать файл: {exc}", file=sys.stderr)
        return 1
    try:
        text = payload.decode("utf-8")
    except UnicodeDecodeError as exc:
        print(f"error: payload не является UTF-8: {exc}", file=sys.stderr)
        return 1
    flag = text.split("\n", 1)[0].strip()
    print(flag)
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))