#!/usr/bin/env python3
"""Black-box adversarial attack against NeuralBadge.

The classifier input is a 32x32 float32 luminance map obtained by:
  crop PATTERN_BOX (96x96) -> grayscale -> 3x3 mean pool -> 32x32.
We control each 3x3 block uniformly, so we effectively control a 32x32
grayscale image with values k/255.  The /batch endpoint returns softmax
probabilities, which let us read the logit margins directly.
"""
from __future__ import annotations

import base64
import io
import json
import sys
import time
from urllib.request import Request, urlopen
from urllib.error import HTTPError

import numpy as np
from PIL import Image

SERVER = "https://peak-privilege.ctf.blackalps.ch"
BADGE_SIZE = (256, 176)          # (W, H)
PATTERN_BOX = (32, 48, 128, 144)  # left, upper, right, lower -> 96x96
MAX_BATCH = 40


def grid_to_png(grid: np.ndarray) -> bytes:
    """grid: 32x32 float in [0,1] -> PNG bytes of a 256x176 RGB badge."""
    g = np.clip(grid, 0.0, 1.0)
    px = np.rint(g * 255.0).astype(np.uint8)            # 32x32 uint8
    block = np.repeat(np.repeat(px, 3, axis=0), 3, axis=1)  # 96x96
    img = np.zeros((BADGE_SIZE[1], BADGE_SIZE[0], 3), dtype=np.uint8)  # H,W,3
    l, u, r, d = PATTERN_BOX
    patch = np.stack([block, block, block], axis=-1)    # 96x96x3
    img[u:d, l:r, :] = patch
    out = io.BytesIO()
    Image.fromarray(img, "RGB").save(out, format="PNG")
    return out.getvalue()


def _post(route: str, body: bytes, content_type: str) -> dict:
    req = Request(SERVER + route, data=body,
                  headers={"Content-Type": content_type}, method="POST")
    for attempt in range(6):
        try:
            with urlopen(req, timeout=60) as resp:
                return json.load(resp)
        except HTTPError as e:
            with e:
                if e.code == 422:
                    return json.load(e)
                if e.code == 503 and attempt < 5:
                    time.sleep(1.5)
                    continue
                raise RuntimeError(f"HTTP {e.code}: {e.read()[:200]}") from e
    raise RuntimeError("retries exhausted")


def query_batch(grids: list[np.ndarray]) -> list[dict]:
    assert 1 <= len(grids) <= MAX_BATCH
    pngs = [grid_to_png(g) for g in grids]
    body = json.dumps({
        "images": [base64.b64encode(p).decode() for p in pngs]
    }).encode()
    return _post("/batch", body, "application/json")["results"]


def eval_grids(grids: list[np.ndarray]) -> list[dict]:
    out = []
    for i in range(0, len(grids), MAX_BATCH):
        out.extend(query_batch(grids[i:i + MAX_BATCH]))
    return out


def margin(res: dict) -> float:
    """log p(ADMIN) - log p(EMPLOYEE); monotone in admin preference."""
    p = res["probabilities"]
    return np.log(max(p["ADMIN"], 1e-300)) - np.log(max(p["EMPLOYEE"], 1e-300))


if __name__ == "__main__":
    # sanity: a few constant grids
    tests = [np.full((32, 32), v, np.float32) for v in (0.0, 0.5, 1.0)]
    for v, r in zip((0.0, 0.5, 1.0), eval_grids(tests)):
        p = r["probabilities"]
        print(f"const {v}: ADMIN={p['ADMIN']:.3e} EMP={p['EMPLOYEE']:.3e} "
              f"INV={p['INVALID']:.3e} access={r['access']}")