$ Forensics: Meow Transmission

Flag: BITSCTF{4rn0ld5_c4t_m4ps_4r3_p3r10d1c}

python3 arnold_lsb_solver.py transmission.png --out . --topk 12

#!/usr/bin/env python3
from __future__ import annotations

import argparse
from pathlib import Path
from typing import List, Tuple

import numpy as np
from PIL import Image

# Challenge parameters from PNG metadata.
N = 128
STAGES = [
    ("m1", 47, 96),
    ("m2", 37, 64),
    ("m1", 29, 96),
]

# Arnold map candidates with order 96 and 64 mod 128.
M1_CANDIDATES = [
    np.array([[1, 1], [1, 2]], dtype=np.int64),
    np.array([[2, 1], [1, 1]], dtype=np.int64),
]

M2_CANDIDATES = [
    np.array([[1, 0], [2, 1]], dtype=np.int64),
    np.array([[1, 1], [2, 3]], dtype=np.int64),
    np.array([[1, 2], [0, 1]], dtype=np.int64),
    np.array([[1, 2], [1, 3]], dtype=np.int64),
    np.array([[2, 1], [3, 2]], dtype=np.int64),
    np.array([[2, 3], [1, 2]], dtype=np.int64),
    np.array([[3, 1], [2, 1]], dtype=np.int64),
    np.array([[3, 2], [1, 1]], dtype=np.int64),
]

rows, cols = np.indices((N, N))

def arnold_step(
    arr: np.ndarray,
    A: np.ndarray,
    coord_mode: int,
    assign_mode: int,
) -> np.ndarray:
    """
    coord_mode:
      0 => x=row, y=col
      1 => x=col, y=row

    assign_mode:
      0 => out[r2,c2] = in[r,c]
      1 => out[r,c]  = in[r2,c2]
    """
    if coord_mode == 0:
        x, y = rows, cols
    else:
        x, y = cols, rows

    xp = (A[0, 0] * x + A[0, 1] * y) % N
    yp = (A[1, 0] * x + A[1, 1] * y) % N

    if coord_mode == 0:
        r2, c2 = xp, yp
    else:
        r2, c2 = yp, xp

    out = np.empty_like(arr)
    if assign_mode == 0:
        out[r2, c2] = arr[rows, cols]
    else:
        out[rows, cols] = arr[r2, c2]
    return out

def arnold_apply(
    arr: np.ndarray,
    A: np.ndarray,
    iterations: int,
    coord_mode: int,
    assign_mode: int,
) -> np.ndarray:
    out = arr
    for _ in range(iterations):
        out = arnold_step(out, A, coord_mode, assign_mode)
    return out

def component_score(bin_img: np.ndarray) -> Tuple[int, float, int]:
    """Return (max_component_size, mean_component_size, num_components)."""
    h, w = bin_img.shape
    vis = np.zeros_like(bin_img, dtype=bool)
    sizes: List[int] = []
    dirs = [(-1, 0), (1, 0), (0, -1), (0, 1)]

    for i in range(h):
        for j in range(w):
            if not bin_img[i, j] or vis[i, j]:
                continue
            stack = [(i, j)]
            vis[i, j] = True
            sz = 0
            while stack:
                x, y = stack.pop()
                sz += 1
                for dx, dy in dirs:
                    nx, ny = x + dx, y + dy
                    if 0 <= nx < h and 0 <= ny < w and bin_img[nx, ny] and not vis[nx, ny]:
                        vis[nx, ny] = True
                        stack.append((nx, ny))
            sizes.append(sz)

    if not sizes:
        return (0, 0.0, 0)
    s = np.array(sizes)
    return (int(s.max()), float(s.mean()), int(len(sizes)))

def solve(input_png: Path, out_dir: Path, topk: int) -> None:
    img = np.array(Image.open(input_png))
    bit = (img & 1).astype(np.uint8)

    results = []
    for i, m1 in enumerate(M1_CANDIDATES):
        for j, m2 in enumerate(M2_CANDIDATES):
            for mode in ("inv", "fwd"):
                for coord_mode in (0, 1):
                    for assign_mode in (0, 1):
                        out = bit.copy()
                        stages = [(m1, 47, 96), (m2, 37, 64), (m1, 29, 96)]

                        # Reverse stage order to decrypt.
                        for A, it, per in stages[::-1]:
                            k = (per - it) if mode == "inv" else it
                            out = arnold_apply(out, A, k, coord_mode, assign_mode)

                        mx, mean, ncomp = component_score(out)
                        results.append(
                            {
                                "mx": mx,
                                "mean": mean,
                                "ncomp": ncomp,
                                "m1_idx": i,
                                "m2_idx": j,
                                "mode": mode,
                                "coord": coord_mode,
                                "assign": assign_mode,
                                "img": out,
                            }
                        )

    # Favor candidates with largest connected text-like strokes.
    results.sort(key=lambda r: (r["mx"], r["mean"], -r["ncomp"]), reverse=True)

    out_dir.mkdir(parents=True, exist_ok=True)
    for rank, r in enumerate(results[:topk]):
        name = (
            f"lsb_rank{rank}_m1{r['m1_idx']}_m2{r['m2_idx']}_"
            f"{r['mode']}_coord{r['coord']}_assign{r['assign']}.png"
        )
        Image.fromarray((r["img"] * 255).astype(np.uint8)).save(out_dir / name)

    # Also save a stable alias for easiest visual inspection.
    Image.fromarray((results[0]["img"] * 255).astype(np.uint8)).save(out_dir / "lsb_rank0.png")

    print("Top candidate:")
    top = results[0]
    print(
        f"  max_comp={top['mx']} mean_comp={top['mean']:.2f} ncomp={top['ncomp']} "
        f"m1={top['m1_idx']} m2={top['m2_idx']} mode={top['mode']} "
        f"coord={top['coord']} assign={top['assign']}"
    )
    print(f"Saved {min(topk, len(results))} ranked images to: {out_dir}")

def main() -> None:
    p = argparse.ArgumentParser(description="Recover LSB text via Arnold-map brute force.")
    p.add_argument("input", type=Path, help="Input PNG (e.g., transmission.png)")
    p.add_argument("--out", type=Path, default=Path("."), help="Output directory")
    p.add_argument("--topk", type=int, default=12, help="How many ranked outputs to save")
    args = p.parse_args()

    solve(args.input, args.out, args.topk)

if __name__ == "__main__":
    main()