roi3.py

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import cv2
import numpy as np
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


ROI_RECURSION_THRESHOLD = 256
MIN_ROI_RECURSION_DEPTH = 1
# Kept only for compatibility with older experiment scripts.
# The FB-RBE core now always uses full-ring scan internally.
ROI_RING_MAX_ROUNDS = 0


RoiBox = tuple[int, int, int, int]


@dataclass
class PipelineResult:
    cipher_image: np.ndarray
    rois: list[RoiBox]
    ownership_map: np.ndarray
    aes_key: bytes
    branch_nonces: dict[str, bytes]
    image_hash: bytes
    roi_keystream_bytes: int
    timings: dict[str, float]


class AESCTRKeyStream:
    """AES-CTR keystream reader with an internal buffer."""

    def __init__(self, key: bytes, nonce: bytes, chunk_size: int = 1 << 20):
        if len(nonce) != 16:
            raise ValueError("AES-CTR nonce must be 16 bytes")
        self._encryptor = Cipher(algorithms.AES(key), modes.CTR(nonce)).encryptor()
        self._chunk_size = int(chunk_size)
        self._buffer = b""
        self._offset = 0
        self.bytes_read = 0

    def _ensure(self, length: int) -> None:
        available = len(self._buffer) - self._offset
        if available >= length:
            return

        remaining = self._buffer[self._offset :] if available > 0 else b""
        need = max(self._chunk_size, length - len(remaining))
        self._buffer = remaining + self._encryptor.update(b"\x00" * need)
        self._offset = 0

    def read(self, length: int) -> bytes:
        if length <= 0:
            return b""
        self._ensure(length)
        start = self._offset
        end = start + length
        self._offset = end
        self.bytes_read += length
        return self._buffer[start:end]

    def read_array(self, length: int) -> np.ndarray:
        return np.frombuffer(self.read(length), dtype=np.uint8)

    def read_u32(self) -> int:
        return int.from_bytes(self.read(4), "little")

    def read_bit(self) -> int:
        return self.read(1)[0] & 1


def normalize_image(image: np.ndarray) -> np.ndarray:
    image = np.asarray(image)
    if image.ndim not in (2, 3):
        raise ValueError("image must be a 2-D grayscale or 3-D color array")
    if image.dtype != np.uint8:
        image = np.clip(image, 0, 255).astype(np.uint8)
    if image.ndim == 3 and image.shape[2] == 4:
        image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
    return np.ascontiguousarray(image)


def load_image(path: str | Path) -> np.ndarray:
    image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
    if image is None:
        raise ValueError(f"Cannot read image: {path}")
    return normalize_image(image)


def save_image(path: str | Path, image: np.ndarray) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    if not cv2.imwrite(str(path), image):
        raise OSError(f"Failed to save image: {path}")


def clamp_roi_box(box: Iterable[int | float], width: int, height: int) -> RoiBox | None:
    x1, y1, x2, y2 = [int(round(value)) for value in box]
    x1 = max(0, min(x1, width))
    x2 = max(0, min(x2, width))
    y1 = max(0, min(y1, height))
    y2 = max(0, min(y2, height))
    if x2 <= x1 or y2 <= y1:
        return None
    return x1, y1, x2, y2


def build_roi_ownership_map(
    image_shape: tuple[int, ...],
    rois: Iterable[Iterable[int | float]],
) -> np.ndarray:
    h, w = image_shape[:2]
    ownership = np.full((h, w), -1, dtype=np.int32)
    for roi_id, box in enumerate(rois):
        clipped = clamp_roi_box(box, w, h)
        if clipped is None:
            continue
        x1, y1, x2, y2 = clipped
        area = ownership[y1:y2, x1:x2]
        area[area < 0] = roi_id
    return ownership


def build_roi_mask(
    image_shape: tuple[int, ...],
    rois: Iterable[Iterable[int | float]],
) -> np.ndarray:
    return build_roi_ownership_map(image_shape, rois) >= 0


def aes_ctr_keystream(key: bytes, nonce: bytes, length: int) -> bytes:
    return AESCTRKeyStream(key, nonce).read(length)


def fisher_yates_permutation(size: int, random_bytes: bytes) -> np.ndarray:
    perm = np.arange(size, dtype=np.int64)
    if size <= 1:
        return perm

    offset = 0
    for i in range(size - 1, 0, -1):
        rnd = int.from_bytes(random_bytes[offset : offset + 4], "little")
        j = rnd % (i + 1)
        perm[i], perm[j] = perm[j], perm[i]
        offset += 4
    return perm


def ring_feistel_like_mix(
    block: np.ndarray,
    stream: AESCTRKeyStream,
    threshold: int = ROI_RECURSION_THRESHOLD,
    min_depth: int = MIN_ROI_RECURSION_DEPTH,
    depth: int = 0,
) -> np.ndarray:
    """
    FB-RBE core with full-ring scan.

    Top layer (depth == 0):
    1. A influences B through ring-tail broadcast.
    2. Swap A and B once.
    3. The swapped A, i.e. the original B, influences the swapped B.
    4. Do not swap again before recursion.

    Deeper layers (depth > 0):
    1. A influences B through one ring-tail broadcast round.
    2. Recurse in swapped order, i.e. Enc(B) || Enc(A).
    """

    def _rotl8(values: np.ndarray, shifts: np.ndarray) -> np.ndarray:
        values16 = values.astype(np.uint16)
        shifts16 = shifts.astype(np.uint16) & np.uint16(7)
        rotated = (
            (values16 << shifts16)
            | (values16 >> ((8 - shifts16) & 7))
        ) & np.uint16(0xFF)
        return rotated.astype(np.uint8)

    def _mix_once(
        a_branch: np.ndarray,
        b_branch: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        """One source-to-target FB-RBE round: A ring-broadcasts to B."""
        a_branch = np.ascontiguousarray(a_branch, dtype=np.uint8)
        b_branch = np.ascontiguousarray(b_branch, dtype=np.uint8)

        p = int(a_branch.size)
        q = int(b_branch.size)
        if p == 0 or q == 0:
            return a_branch, b_branch

        direction = stream.read_bit()
        step = 1 + (stream.read_u32() % (p - 1)) if p > 1 else 0
        if p > 1 and step != 0:
            while int(np.gcd(step, p)) != 1:
                step += 1
                if step >= p:
                    step = 1

        # Full-ring scan: the number of virtual ring moves is always at least
        # the source branch length, and the exact extra scan is keystream-driven.
        rounds = p + (stream.read_u32() % p) if p > 1 else 1

        if p > 1 and step != 0:
            round_ids = np.arange(1, rounds + 1, dtype=np.int64)
            offsets = (round_ids * int(step)) % p
            if direction == 0:
                tail_indices = (p - 1 - offsets) % p
            else:
                tail_indices = offsets
            rotation_amount = int((rounds * int(step)) % p)
        else:
            tail_indices = np.zeros(rounds, dtype=np.int64)
            rotation_amount = 0

        # Collect many virtual tails and aggregate them into a 16-bit state.
        tail_values = a_branch[tail_indices].astype(np.uint64)
        tail_state = int(np.sum(tail_values, dtype=np.uint64) & np.uint64(0xFFFF))
        tail_low = np.uint16(tail_state & 0xFF)
        tail_high = np.uint16((tail_state >> 8) & 0xFF)

        b_positions = np.arange(q, dtype=np.uint16) & np.uint16(0xFF)
        pos_mul = np.uint16(int(stream.read(1)[0]) | 1)
        pos_add = np.uint16(int(stream.read(1)[0]))
        g_key = stream.read_array(q).astype(np.uint16)
        rot_key = (
            stream.read_array(q).astype(np.uint16)
            + tail_low
            + tail_high
            + b_positions
        ) & np.uint16(7)

        tail_term = (
            tail_low
            + ((b_positions + np.uint16(1)) * tail_high)
        ) & np.uint16(0xFF)
        pos_term = (b_positions * pos_mul + pos_add) & np.uint16(0xFF)
        g_value = (g_key + pos_term + tail_term) & np.uint16(0xFF)

        b_added = (b_branch.astype(np.uint16) + g_value) & np.uint16(0xFF)
        b_branch = _rotl8(b_added.astype(np.uint8), rot_key.astype(np.uint8))

        # Apply the accumulated rotation to A only once.
        if p > 1 and rotation_amount != 0:
            if direction == 0:
                a_branch = np.concatenate(
                    (a_branch[-rotation_amount:], a_branch[:-rotation_amount])
                ).astype(np.uint8, copy=False)
            else:
                a_branch = np.concatenate(
                    (a_branch[rotation_amount:], a_branch[:rotation_amount])
                ).astype(np.uint8, copy=False)
        else:
            a_branch = a_branch.copy()

        # Lightweight branch whitening: XOR is enough here because the broadcast
        # stage above already contains modular addition and ROTL8.
        a_key = stream.read_array(p)
        b_key = stream.read_array(q)
        a_branch = np.bitwise_xor(a_branch, a_key).astype(np.uint8, copy=False)
        b_branch = np.bitwise_xor(b_branch, b_key).astype(np.uint8, copy=False)
        return a_branch, b_branch

    block = np.ascontiguousarray(block, dtype=np.uint8)
    n = int(block.size)

    if n == 0:
        return block.copy()

    if n <= 1 or (n <= threshold and depth >= min_depth):
        leaf_key = stream.read_array(n)
        return np.bitwise_xor(block, leaf_key).astype(np.uint8, copy=False)

    split = n // 2
    a_branch = np.ascontiguousarray(block[:split], dtype=np.uint8)
    b_branch = np.ascontiguousarray(block[split:], dtype=np.uint8)

    if a_branch.size == 0 or b_branch.size == 0:
        leaf_key = stream.read_array(n)
        return np.bitwise_xor(block, leaf_key).astype(np.uint8, copy=False)

    if depth == 0:
        # First top-level pass: original A affects original B.
        a_branch, b_branch = _mix_once(a_branch, b_branch)

        # Swap once, then let the original B side affect the original A side.
        a_branch, b_branch = b_branch, a_branch
        a_branch, b_branch = _mix_once(a_branch, b_branch)

        # No extra top-level swap after the second pass.
        left_enc = ring_feistel_like_mix(
            a_branch,
            stream,
            threshold,
            min_depth,
            depth + 1,
        )
        right_enc = ring_feistel_like_mix(
            b_branch,
            stream,
            threshold,
            min_depth,
            depth + 1,
        )
    else:
        # Deeper layers keep the original one-pass-then-swap recursion rule.
        a_branch, b_branch = _mix_once(a_branch, b_branch)
        left_enc = ring_feistel_like_mix(
            b_branch,
            stream,
            threshold,
            min_depth,
            depth + 1,
        )
        right_enc = ring_feistel_like_mix(
            a_branch,
            stream,
            threshold,
            min_depth,
            depth + 1,
        )

    return np.concatenate((left_enc, right_enc)).astype(np.uint8, copy=False)

def roi_encryptor(
    image: np.ndarray,
    ownership_map: np.ndarray,
    roi_stream: AESCTRKeyStream,
    threshold: int = ROI_RECURSION_THRESHOLD,
    min_depth: int = MIN_ROI_RECURSION_DEPTH,
) -> np.ndarray:
    rows, cols = np.nonzero(ownership_map >= 0)
    if rows.size == 0:
        return image

    roi_vector = np.ascontiguousarray(image[rows, cols])
    roi_bytes = np.ascontiguousarray(roi_vector.reshape(-1), dtype=np.uint8)
    if roi_bytes.size == 0:
        return image

    encrypted_bytes = ring_feistel_like_mix(
        roi_bytes,
        roi_stream,
        threshold=threshold,
        min_depth=min_depth,
    )
    image[rows, cols] = encrypted_bytes.reshape(roi_vector.shape)
    return image


def derive_branch_nonce(label: bytes, shape_bytes: bytes, image_hash: bytes) -> bytes:
    return hashlib.sha256(
        b"roi-aes-ctr-nonce-v3|" + label + b"|" + shape_bytes + b"|" + image_hash
    ).digest()[:16]


def keyderiver(
    masterkey: str,
    width: int,
    height: int,
    channels: int = 3,
    image_hash: bytes | None = None,
) -> tuple[bytes, dict[str, bytes]]:
    if image_hash is None:
        image_hash = hashlib.sha256(
            f"{width}x{height}x{channels}".encode("ascii")
        ).digest()
    shape_bytes = f"{width}x{height}x{channels}".encode("ascii")
    key_salt = hashlib.sha256(
        b"roi-aes-key-salt-v3|" + shape_bytes + b"|" + image_hash
    ).digest()[:16]
    derivation = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=key_salt,
        iterations=240000,
    )
    aes_key = derivation.derive(masterkey.encode("utf-8"))
    nonces = {
        "roi": derive_branch_nonce(b"roi", shape_bytes, image_hash),
        "global_diffusion": derive_branch_nonce(
            b"global-diffusion", shape_bytes, image_hash
        ),
        "global_row": derive_branch_nonce(b"global-row", shape_bytes, image_hash),
        "global_col": derive_branch_nonce(b"global-col", shape_bytes, image_hash),
    }
    return aes_key, nonces


def load_yolo_model(model_path: str | Path | None = None):
    from ultralytics import YOLO

    if model_path is None:
        local_model = Path(__file__).with_name("yolo11n.pt")
        model_path = local_model if local_model.exists() else "yolov8n.pt"
    return YOLO(str(model_path))


def detector(
    model,
    image: np.ndarray,
    confidence: float,
    classes: set[int] | None = None,
) -> list[RoiBox]:
    h, w = image.shape[:2]
    detector_image = image
    if image.ndim == 2:
        detector_image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    elif image.ndim == 3 and image.shape[2] == 4:
        detector_image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)

    result = model.predict(source=detector_image, conf=confidence, verbose=False)[0]
    if result.boxes is None or len(result.boxes) == 0:
        return []

    boxes = result.boxes.xyxy.cpu().numpy().astype(int)
    confs = result.boxes.conf.cpu().numpy()
    cls_ids = result.boxes.cls.cpu().numpy().astype(int)
    rois: list[RoiBox] = []
    for box, score, clsid in sorted(
        zip(boxes, confs, cls_ids),
        key=lambda item: item[1],
        reverse=True,
    ):
        if classes is not None and clsid not in classes:
            continue
        clipped = clamp_roi_box(box.tolist(), w, h)
        if clipped is not None:
            rois.append(clipped)
    return rois


def diffusion(
    image: np.ndarray,
    keystream: bytes,
    x1: int,
    y1: int,
    x2: int,
    y2: int,
) -> np.ndarray:
    roi = image[y1:y2, x1:x2]
    bytes_needed = roi.size
    if len(keystream) < bytes_needed or bytes_needed == 0:
        return image
    key_arr = np.frombuffer(keystream[:bytes_needed], dtype=np.uint8).reshape(
        roi.shape
    )
    image[y1:y2, x1:x2] = np.bitwise_xor(roi, key_arr)
    return image


def permutation(
    image: np.ndarray,
    col_keystream: bytes,
    row_keystream: bytes,
) -> np.ndarray:
    h, w = image.shape[:2]
    row_perm = fisher_yates_permutation(h, row_keystream)
    col_perm = fisher_yates_permutation(w, col_keystream)

    tmp = np.empty_like(image)
    np.take(image, row_perm, axis=0, out=tmp)
    np.take(tmp, col_perm, axis=1, out=image)
    return image


def encrypt_image(
    image: np.ndarray,
    masterkey: str,
    rois: Iterable[Iterable[int | float]] | None = None,
    *,
    yolo_model=None,
    confidence: float = 0.25,
    classes: set[int] | None = None,
    use_image_hash: bool = True,
    copy: bool = True,
    threshold: int = ROI_RECURSION_THRESHOLD,
    min_depth: int = MIN_ROI_RECURSION_DEPTH,
) -> PipelineResult:
    """Run the roi2 encryption pipeline without any GUI or file side effects.

    If rois is None and yolo_model is provided, ROIs are detected with YOLO.
    If rois is None and yolo_model is also None, ROI encryption is skipped and
    only the full-image row/column permutation plus diffusion stages run.
    """

    timings: dict[str, float] = {}
    total_start = time.perf_counter()

    stage_start = time.perf_counter()
    source = normalize_image(image)
    cipher_image = source.copy() if copy else source
    timings["copy_image"] = time.perf_counter() - stage_start

    h, w = cipher_image.shape[:2]
    channels = 1 if cipher_image.ndim == 2 else cipher_image.shape[2]
    image_hash = (
        hashlib.sha256(np.ascontiguousarray(cipher_image).tobytes()).digest()
        if use_image_hash
        else hashlib.sha256(f"{w}x{h}x{channels}".encode("ascii")).digest()
    )

    stage_start = time.perf_counter()
    aes_key, branch_nonces = keyderiver(masterkey, w, h, channels, image_hash)
    timings["key_derivation"] = time.perf_counter() - stage_start

    stage_start = time.perf_counter()
    if rois is None:
        detected_rois = (
            detector(yolo_model, cipher_image, confidence, classes)
            if yolo_model is not None
            else []
        )
    else:
        detected_rois = []
        for box in rois:
            clipped = clamp_roi_box(box, w, h)
            if clipped is not None:
                detected_rois.append(clipped)
    timings["roi_detection_or_selection"] = time.perf_counter() - stage_start

    stage_start = time.perf_counter()
    ownership_map = build_roi_ownership_map(cipher_image.shape, detected_rois)
    timings["roi_ownership_map"] = time.perf_counter() - stage_start

    stage_start = time.perf_counter()
    roi_keystream_bytes = 0
    if np.any(ownership_map >= 0):
        roi_stream = AESCTRKeyStream(aes_key, branch_nonces["roi"])
        roi_encryptor(
            cipher_image,
            ownership_map,
            roi_stream,
            threshold=threshold,
            min_depth=min_depth,
        )
        roi_keystream_bytes = int(roi_stream.bytes_read)
    timings["roi_feistel_encrypt"] = time.perf_counter() - stage_start

    stage_start = time.perf_counter()
    row_random = aes_ctr_keystream(
        aes_key,
        branch_nonces["global_row"],
        max(0, h - 1) * 4,
    )
    col_random = aes_ctr_keystream(
        aes_key,
        branch_nonces["global_col"],
        max(0, w - 1) * 4,
    )
    permutation(cipher_image, col_random, row_random)
    timings["global_row_col_permutation"] = time.perf_counter() - stage_start

    stage_start = time.perf_counter()
    global_diffusion_key = aes_ctr_keystream(
        aes_key,
        branch_nonces["global_diffusion"],
        cipher_image.size,
    )
    diffusion(cipher_image, global_diffusion_key, 0, 0, w, h)
    timings["global_diffusion"] = time.perf_counter() - stage_start

    timings["total"] = time.perf_counter() - total_start
    return PipelineResult(
        cipher_image=cipher_image,
        rois=detected_rois,
        ownership_map=ownership_map,
        aes_key=aes_key,
        branch_nonces=branch_nonces,
        image_hash=image_hash,
        roi_keystream_bytes=roi_keystream_bytes,
        timings=timings,
    )