240 lines
9.6 KiB
Python
240 lines
9.6 KiB
Python
"""Lossless carrier codec used by ChickHideNode.
|
|
|
|
The carrier is a lossless RGB PNG. A four-byte big-endian header length is
|
|
stored first, followed by a binary file header and its payload. The header
|
|
contains an optional password marker, extension and payload length.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from hashlib import sha256
|
|
from hmac import compare_digest
|
|
from io import BytesIO
|
|
from numbers import Integral
|
|
from pathlib import Path
|
|
import os
|
|
import struct
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
SKIP_W_RATIO = 0.40
|
|
SKIP_H_RATIO = 0.08
|
|
CHANNELS = 3
|
|
MINIMUM_SIDE = 640
|
|
SUPPORTED_LSB_BITS = (2, 6, 8)
|
|
|
|
|
|
class CarrierCodecError(ValueError):
|
|
"""Raised when a carrier cannot be encoded or decoded."""
|
|
|
|
|
|
def _key_stream(password: str, salt: bytes, length: int) -> bytes:
|
|
material = (password + salt.hex()).encode("utf-8")
|
|
output = bytearray()
|
|
counter = 0
|
|
while len(output) < length:
|
|
output.extend(sha256(material + str(counter).encode("ascii")).digest())
|
|
counter += 1
|
|
return bytes(output[:length])
|
|
|
|
|
|
def _protect(raw: bytes, password: str) -> tuple[bytes, bytes, bytes, bool]:
|
|
if not password:
|
|
return raw, b"", b"", False
|
|
salt = os.urandom(16)
|
|
check = sha256((password + salt.hex()).encode("utf-8")).digest()
|
|
stream = _key_stream(password, salt, len(raw))
|
|
return bytes(left ^ right for left, right in zip(raw, stream)), salt, check, True
|
|
|
|
|
|
def build_file_header(raw: bytes, password: str, extension: str) -> bytes:
|
|
extension = extension.lstrip(".")
|
|
extension_bytes = extension.encode("utf-8")
|
|
if not extension_bytes or len(extension_bytes) > 255:
|
|
raise CarrierCodecError("extension must contain 1-255 UTF-8 bytes")
|
|
protected, salt, check, has_password = _protect(raw, password)
|
|
header = bytearray((1 if has_password else 0,))
|
|
if has_password:
|
|
header.extend(check)
|
|
header.extend(salt)
|
|
header.append(len(extension_bytes))
|
|
header.extend(extension_bytes)
|
|
header.extend(struct.pack(">I", len(protected)))
|
|
header.extend(protected)
|
|
return bytes(header)
|
|
|
|
|
|
def parse_file_header(header: bytes, password: str = "") -> tuple[bytes, str]:
|
|
if len(header) < 1:
|
|
raise CarrierCodecError("carrier header is empty")
|
|
cursor = 1
|
|
has_password = header[0] == 1
|
|
if header[0] not in (0, 1):
|
|
raise CarrierCodecError("carrier password marker is invalid")
|
|
check = salt = b""
|
|
if has_password:
|
|
if len(header) < cursor + 48:
|
|
raise CarrierCodecError("carrier password header is truncated")
|
|
check = header[cursor:cursor + 32]
|
|
salt = header[cursor + 32:cursor + 48]
|
|
cursor += 48
|
|
if len(header) < cursor + 1:
|
|
raise CarrierCodecError("carrier extension is missing")
|
|
extension_length = header[cursor]
|
|
cursor += 1
|
|
if len(header) < cursor + extension_length + 4:
|
|
raise CarrierCodecError("carrier metadata is truncated")
|
|
try:
|
|
extension = header[cursor:cursor + extension_length].decode("utf-8")
|
|
except UnicodeDecodeError as exc:
|
|
raise CarrierCodecError("carrier extension is invalid") from exc
|
|
cursor += extension_length
|
|
payload_length = struct.unpack(">I", header[cursor:cursor + 4])[0]
|
|
cursor += 4
|
|
payload = header[cursor:]
|
|
if len(payload) != payload_length:
|
|
raise CarrierCodecError("carrier payload length does not match")
|
|
if not has_password:
|
|
return payload, extension
|
|
if not password:
|
|
raise CarrierCodecError("password is required")
|
|
expected = sha256((password + salt.hex()).encode("utf-8")).digest()
|
|
if not compare_digest(expected, check):
|
|
raise CarrierCodecError("password is incorrect")
|
|
stream = _key_stream(password, salt, len(payload))
|
|
return bytes(left ^ right for left, right in zip(payload, stream)), extension
|
|
|
|
|
|
def _lsb_bits(compress: int) -> int:
|
|
if isinstance(compress, bool) or not isinstance(compress, Integral) or compress not in SUPPORTED_LSB_BITS:
|
|
raise CarrierCodecError("compress must be 2, 6 or 8")
|
|
return int(compress)
|
|
|
|
|
|
def _channel_count(height: int, width: int) -> int:
|
|
return (height * width - int(height * SKIP_H_RATIO) * int(width * SKIP_W_RATIO)) * CHANNELS
|
|
|
|
|
|
def _channel_chunks(array):
|
|
"""Yield writable row-major views, omitting the reserved rectangle."""
|
|
height, width, _ = array.shape
|
|
skip_h, skip_w = int(height * SKIP_H_RATIO), int(width * SKIP_W_RATIO)
|
|
flat = array.reshape(-1)
|
|
spans = ((row * width * 3 + skip_w * 3, (row + 1) * width * 3) for row in range(skip_h))
|
|
for start, stop in spans:
|
|
yield flat[start:stop]
|
|
for start in range(skip_h * width * 3, len(flat), 65536):
|
|
yield flat[start:start + 65536]
|
|
|
|
|
|
def required_side(header_length: int, compress: int, minimum: int = MINIMUM_SIDE) -> int:
|
|
bits = _lsb_bits(compress)
|
|
side = max(MINIMUM_SIDE, int(minimum))
|
|
needed = (header_length + 4) * 8
|
|
while _channel_count(side, side) * bits < needed:
|
|
side += 64
|
|
return side
|
|
|
|
|
|
def _embed(canvas: Image.Image, header: bytes, compress: int) -> Image.Image:
|
|
bits_per_channel = _lsb_bits(compress)
|
|
payload = struct.pack(">I", len(header)) + header
|
|
array = np.array(canvas.convert("RGB"), dtype=np.uint8)
|
|
if len(payload) * 8 > _channel_count(*array.shape[:2]) * bits_per_channel:
|
|
raise CarrierCodecError("carrier capacity is too small")
|
|
keep_mask = np.uint8(255 ^ ((1 << bits_per_channel) - 1))
|
|
weights = 1 << np.arange(bits_per_channel - 1, -1, -1, dtype=np.uint8)
|
|
position = 0
|
|
for channels in _channel_chunks(array):
|
|
count = min(len(channels), (len(payload) * 8 - position + bits_per_channel - 1) // bits_per_channel)
|
|
if count <= 0:
|
|
break
|
|
end = min(position + count * bits_per_channel, len(payload) * 8)
|
|
bits = np.unpackbits(np.frombuffer(payload[position // 8:(end + 7) // 8], dtype=np.uint8))
|
|
bits = bits[position % 8:position % 8 + end - position]
|
|
bits = np.pad(bits, (0, count * bits_per_channel - len(bits)))
|
|
values = (bits.reshape(count, bits_per_channel) * weights).sum(axis=1).astype(np.uint8)
|
|
channels[:count] = (channels[:count] & keep_mask) | values
|
|
position += count * bits_per_channel
|
|
return Image.fromarray(array)
|
|
|
|
|
|
def extract_file_header(image: Image.Image, compress: int) -> bytes:
|
|
bits_per_channel = _lsb_bits(compress)
|
|
array = np.array(image.convert("RGB"), dtype=np.uint8)
|
|
output, pending = bytearray(), np.empty(0, dtype=np.uint8)
|
|
total_bytes = None
|
|
for channels in _channel_chunks(array):
|
|
values = channels & ((1 << bits_per_channel) - 1)
|
|
bits = np.unpackbits(values).reshape(-1, 8)[:, -bits_per_channel:].reshape(-1)
|
|
bits = np.concatenate((pending, bits))
|
|
end = len(bits) // 8 * 8
|
|
output.extend(np.packbits(bits[:end]).tobytes())
|
|
pending = bits[end:]
|
|
if total_bytes is None and len(output) >= 4:
|
|
total_bytes = 4 + struct.unpack(">I", output[:4])[0]
|
|
if total_bytes <= 4 or total_bytes * 8 > _channel_count(*array.shape[:2]) * bits_per_channel:
|
|
raise CarrierCodecError("carrier header length is invalid")
|
|
if total_bytes is not None and len(output) >= total_bytes:
|
|
return bytes(output[4:total_bytes])
|
|
raise CarrierCodecError("carrier is too small")
|
|
|
|
|
|
def decode_bytes(image: Image.Image, password: str = "") -> tuple[bytes, str]:
|
|
errors = []
|
|
for compress in SUPPORTED_LSB_BITS:
|
|
try:
|
|
return parse_file_header(extract_file_header(image, compress), password)
|
|
except CarrierCodecError as exc:
|
|
errors.append(str(exc))
|
|
raise CarrierCodecError("carrier decode failed: " + "; ".join(errors))
|
|
|
|
|
|
def _make_canvas(side: int, title: str) -> Image.Image:
|
|
image = Image.new("RGB", (side, side), (255, 248, 233))
|
|
with Image.open(Path(__file__).resolve().parents[1] / "assets" / "chick.png") as source:
|
|
cover = source.convert("RGBA").resize((side, side), Image.Resampling.LANCZOS)
|
|
image.paste(cover, (0, 0), cover.getchannel("A"))
|
|
draw = ImageDraw.Draw(image)
|
|
if title:
|
|
font = ImageFont.load_default()
|
|
draw.text((int(side * .05), int(side * .04)), title[:40], fill=(20, 40, 50), font=font)
|
|
return image
|
|
|
|
|
|
def encode_bytes(raw: bytes, extension: str, password: str = "", compress: int = 2,
|
|
title: str = "", fixed_side: Optional[int] = None) -> Image.Image:
|
|
header = build_file_header(raw, password, extension)
|
|
side = required_side(len(header), compress, fixed_side or MINIMUM_SIDE)
|
|
return _embed(_make_canvas(side, title), header, compress)
|
|
|
|
|
|
def image_pixels(image: object) -> np.ndarray:
|
|
if hasattr(image, "detach"):
|
|
image = image.detach().cpu()
|
|
image = (image.float() if image.is_floating_point() else image).numpy()
|
|
array = np.asarray(image)
|
|
if array.ndim == 4:
|
|
if array.shape[0] != 1:
|
|
raise CarrierCodecError("image_to_png_bytes accepts one frame")
|
|
array = array[0]
|
|
if array.ndim != 3 or array.shape[-1] not in (3, 4) or min(array.shape) < 1:
|
|
raise CarrierCodecError("images must have shape HxWx3 or HxWx4")
|
|
if array.dtype.kind not in "fiu" or not np.isfinite(array).all():
|
|
raise CarrierCodecError("image pixels must be finite numbers")
|
|
if np.issubdtype(array.dtype, np.floating):
|
|
array = np.rint(np.clip(array, 0.0, 1.0) * 255.0).astype(np.uint8)
|
|
else:
|
|
array = np.clip(array, 0, 255).astype(np.uint8)
|
|
return array
|
|
|
|
|
|
def image_to_png_bytes(image: object) -> bytes:
|
|
image = Image.fromarray(image_pixels(image))
|
|
output = BytesIO()
|
|
image.save(output, format="PNG", optimize=True, compress_level=9)
|
|
return output.getvalue()
|