encoder.py

"""
Secret file encoder.
"""

import pathlib
from typing import Iterator

from .bitmap import Bitmap


class EncodingError(Exception):
    pass


class SecretFile:
    """Convenience class for serializing secret data."""

    def __init__(self, path: pathlib.Path):
        self.path = path
        self.filename = path.name.encode("utf-8") + b"\x00"
        self.size_bytes = path.stat().st_size

    @property
    def num_secret_bytes(self) -> int:
        """Total number of bytes including the null-terminated string."""
        return len(self.filename) + self.size_bytes

    @property
    def secret_bytes(self) -> Iterator[int]:
        """Null-terminated name followed by the file content."""
        yield from self.filename
        with self.path.open(mode="rb") as file:
            yield from file.read()


def encode(bitmap: Bitmap, path: pathlib.Path) -> None:
    """Embed a secret file in the bitmap."""

    file = SecretFile(path)

    if file.num_secret_bytes > bitmap.max_bytes:
        raise EncodingError("Not enough pixels to embed a secret file")

    bitmap.reserved_field = file.size_bytes
    for secret_byte, eight_bytes in zip(file.secret_bytes, bitmap.byte_slices):
        secret_bits = [(secret_byte >> i) & 1 for i in reversed(range(8))]
        bitmap[eight_bytes] = bytes(
            [
                byte | 1 if bit else byte & ~1
                for byte, bit in zip(bitmap[eight_bytes], secret_bits)
            ]
        )

    print("Secret file was embedded in the bitmap")