"""

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")