#!/usr/bin/env python3
"""
afterberry_ref.py -- reference implementation, Afterberry Vault Manual revision 2.0

Profiles implemented
  AB-L1  the lock: repeated squaring modulo n, two doors, waypoints, relay proofs
  AB-E1  capsule chunk encryption (AES-256-GCM) and key wrap (RFC 3394)
  AB-S1  the 60-of-88 erasure code that turns one capsule into 88 shards
  AB-R1  the 88-byte photographic record and its RS(128,88) protection
  AB-C1  the photographic carrier: keyed LSB matching in 8-bit RGB PNG

Python 3.8+, standard library only. The two AES operations (key unwrap and chunk
decryption) use the 'cryptography' package when it is installed and are skipped,
with a message, when it is not. Everything else -- the lock, both doors, relay
proofs, the record, Reed-Solomon coding, PNG reading and writing, embedding and
extraction -- needs nothing but this file.

Written to be read. Where speed and clarity disagreed, clarity won.

  python3 afterberry_ref.py lock-verify rehearsal-lock.json --door capability
  python3 afterberry_ref.py lock-verify rehearsal-lock.json --door proofs
  python3 afterberry_ref.py lock-verify rehearsal-lock.json --door patience
  python3 afterberry_ref.py carrier-extract rehearsal-carrier.png --index 0
  python3 afterberry_ref.py record-parse <256 hex characters>
  python3 afterberry_ref.py selftest

Licence: GPL-3.0-or-later. No warranty; see /manual/licence.
"""

import argparse
import hashlib
import hmac
import json
import struct
import sys
import time
import zlib

# --------------------------------------------------------------------------
# integers and bytes
# --------------------------------------------------------------------------

def i2osp(v: int, length: int) -> bytes:
    return v.to_bytes(length, "big")


def os2ip(b: bytes) -> int:
    return int.from_bytes(b, "big")


def sha256(b: bytes) -> bytes:
    return hashlib.sha256(b).digest()


def hkdf_sha256(ikm: bytes, salt: bytes, info: bytes, length: int = 32) -> bytes:
    """RFC 5869 with HMAC-SHA-256."""
    prk = hmac.new(salt, ikm, hashlib.sha256).digest()
    okm, block, counter = b"", b"", 1
    while len(okm) < length:
        block = hmac.new(prk, block + info + bytes([counter]), hashlib.sha256).digest()
        okm += block
        counter += 1
    return okm[:length]


# --------------------------------------------------------------------------
# AB-L1 -- the lock
# --------------------------------------------------------------------------

class Lock:
    """Public lock parameters, as read from a lock file."""

    def __init__(self, doc: dict):
        self.lock_id = doc["lock_id"]
        self.bits = doc["bits"]
        self.n = int(doc["n_hex"], 16)
        # exponents are decimal strings: the real ones do not fit in a float or an int64
        self.waypoints = [(w["ring"], int(w["t"]), bytes.fromhex(w["check_sha256"]))
                          for w in doc["waypoints"]]
        stones = doc.get("milestones", {})
        self.milestone_interval = int(stones.get("interval", "0"))
        self.milestones = [bytes.fromhex(h) for h in stones.get("checks_sha256", [])]
        if stones and (not self.milestones
                       or self.milestone_interval != self.waypoints[-1][1] // len(self.milestones)):
            raise ValueError("milestone interval must be the last exponent divided by the number of milestones")
        self.nbytes = (self.bits + 7) // 8
        if self.n.bit_length() != self.bits:
            raise ValueError("modulus length does not match declared bits")
        self.digest = sha256(self.core())
        if "lock_digest_sha256" in doc and bytes.fromhex(doc["lock_digest_sha256"]) != self.digest:
            raise ValueError("lock digest mismatch")
        self.x = base_element(self.lock_id, self.n, self.nbytes)
        if "x_hex" in doc and int(doc["x_hex"], 16) != self.x:
            raise ValueError("base element mismatch")

    def core(self) -> bytes:
        """Canonical byte string of the lock parameters. Its SHA-256 is the lock digest."""
        out = b"AB8L" + bytes([1]) + struct.pack(">H", self.bits) + i2osp(self.n, self.nbytes)
        out += bytes([len(self.waypoints)])
        for _, t, _ in self.waypoints:
            out += i2osp(t, 16)
        ident = self.lock_id.encode("ascii")
        return out + bytes([len(ident)]) + ident

    def canon(self, z: int) -> int:
        """Canonical representative of {z, n - z}."""
        z %= self.n
        return min(z, self.n - z)

    def check(self, ring: int, y: int) -> bytes:
        return sha256(b"afterberry/AB-L1/check/" + bytes([ring]) + i2osp(y, self.nbytes))

    def milestone(self, j: int, z: int) -> bytes:
        """Check for milestone j (1-based): the chain value after j * interval squarings."""
        return sha256(b"afterberry/AB-L1/milestone/" + struct.pack(">H", j) + i2osp(self.canon(z), self.nbytes))

    def ring_key(self, ring: int, y: int) -> bytes:
        return hkdf_sha256(i2osp(y, self.nbytes), self.digest,
                           b"afterberry/AB-L1/ring/" + bytes([ring]))

    def capsule_kek(self, ring_key: bytes, capsule_id: int) -> bytes:
        return hkdf_sha256(ring_key, self.digest,
                           b"afterberry/AB-L1/capsule/" + struct.pack(">H", capsule_id))


def base_element(lock_id: str, n: int, nbytes: int) -> int:
    """x = <h^2 mod n>, h = SHAKE256(label, nbytes + 16) reduced modulo n."""
    label = b"afterberry/AB-L1/base/" + lock_id.encode("ascii")
    h = os2ip(hashlib.shake_256(label).digest(nbytes + 16)) % n
    z = h * h % n
    return min(z, n - z)


def door_patience(lock: Lock, progress=None):
    """The patient door: t sequential squarings. Yields (ring, t, y) at each waypoint.

    Milestones are checked as they are passed, so a wrong turn is caught within one
    interval instead of at the next waypoint."""
    z, done, step = lock.x, 0, lock.milestone_interval
    for ring, t, _ in lock.waypoints:
        while done < t:
            z = z * z % lock.n
            done += 1
            if step and done % step == 0 and done // step <= len(lock.milestones):
                if lock.milestone(done // step, z) != lock.milestones[done // step - 1]:
                    raise ValueError("milestone %d does not match" % (done // step))
            if progress and done % 250000 == 0:
                progress(done)
        yield ring, t, lock.canon(z)


def door_capability(lock: Lock, p: int, q: int):
    """The capability door: with the factors, each waypoint is one short exponentiation."""
    if p * q != lock.n:
        raise ValueError("p * q != n")
    phi = (p - 1) * (q - 1)
    for ring, t, _ in lock.waypoints:
        e = pow(2, t, phi)
        yield ring, t, lock.canon(pow(lock.x, e, lock.n))


# -- relay proofs (Wesolowski 2019) ----------------------------------------

_SMALL_PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71,
                 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151,
                 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, 223, 227, 229, 233,
                 239, 241, 251, 257, 263, 269, 271, 277, 281, 283, 293, 307, 311]


def is_probable_prime(m: int) -> bool:
    """Miller-Rabin to the first 64 prime bases.

    Fixed bases are adequate here only because every candidate is a hash output.
    Do not use this on numbers chosen by an adversary."""
    if m < 2:
        return False
    for sp in _SMALL_PRIMES:
        if m % sp == 0:
            return m == sp
    d, r = m - 1, 0
    while d % 2 == 0:
        d //= 2
        r += 1
    for a in _SMALL_PRIMES:
        v = pow(a, d, m)
        if v in (1, m - 1):
            continue
        for _ in range(r - 1):
            v = v * v % m
            if v == m - 1:
                break
        else:
            return False
    return True


def hash_to_prime(lock: Lock, a: int, b: int, T: int) -> int:
    seed = (b"afterberry/AB-L1/relay/" + lock.digest + i2osp(a, lock.nbytes)
            + i2osp(b, lock.nbytes) + i2osp(T, 16))
    i = 0
    while True:
        c = os2ip(sha256(seed + struct.pack(">I", i))) | (1 << 255) | 1
        if is_probable_prime(c):
            return c
        i += 1


def relay_prove(lock: Lock, a: int, b: int, T: int) -> int:
    """Proof that b = <a^(2^T)>. Costs about as much as the leg itself."""
    ell = hash_to_prime(lock, a, b, T)
    return lock.canon(pow(a, (1 << T) // ell, lock.n))


def relay_verify(lock: Lock, a: int, b: int, T: int, pi: int) -> bool:
    """Sound against anyone who does not hold the factors of n. Someone who does can
    forge a proof -- and has no need to, since they can open the lock outright."""
    import math
    for v in (a, b, pi):                          # zero, n, or anything sharing a factor with n
        if not 0 < v < lock.n or math.gcd(v, lock.n) != 1:
            return False
    ell = hash_to_prime(lock, a, b, T)
    r = pow(2, T, ell)
    return lock.canon(pow(pi, ell, lock.n) * pow(a, r, lock.n)) == lock.canon(b)


# --------------------------------------------------------------------------
# AB-E1 -- capsule header, key wrap, chunk encryption
# --------------------------------------------------------------------------

CHUNK = 1 << 20
REMAINDER_RING = 0xFF


def header_core(ring, capsule_id, k, n, plaintext_len, lock_digest, chunk=CHUNK) -> bytes:
    core = (b"AB8C" + bytes([1, ring]) + struct.pack(">H", capsule_id) + bytes([k, n, 0, 0])
            + struct.pack(">I", chunk) + struct.pack(">Q", plaintext_len) + lock_digest + bytes(8))
    assert len(core) == 64
    return core


def chunk_nonce(index: int, final: bool) -> bytes:
    return bytes(3) + bytes([1 if final else 0]) + struct.pack(">Q", index)


def _aes():
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        from cryptography.hazmat.primitives.keywrap import aes_key_unwrap, aes_key_wrap
    except ImportError:
        raise RuntimeError("AES steps need the 'cryptography' package (pip install cryptography)")
    return AESGCM, aes_key_wrap, aes_key_unwrap


def seal_capsule(plain: bytes, content_key: bytes, kek, core: bytes) -> bytes:
    """Capsule file = header core || sealed key || chunk records. kek=None makes a Remainder
    capsule: the key field is zero and the content key is simply never kept."""
    AESGCM, wrap, _ = _aes()
    size = struct.unpack(">I", core[12:16])[0]
    assert struct.unpack(">Q", core[16:24])[0] == len(plain)
    count = max(1, -(-len(plain) // size))            # an empty capsule is one empty record
    body = b"".join(AESGCM(content_key).encrypt(chunk_nonce(i, i == count - 1), plain[i * size:(i + 1) * size], core)
                    for i in range(count))
    return core + (wrap(kek, content_key) if kek else bytes(40)) + body


def open_capsule(capsule: bytes, kek: bytes) -> bytes:
    """capsule = header core (64) || wrapped key (40) || chunk records (ciphertext || tag)."""
    AESGCM, _, aes_key_unwrap = _aes()
    core, wrapped, body = capsule[:64], capsule[64:104], capsule[104:]
    size = struct.unpack(">I", core[12:16])[0]
    total = struct.unpack(">Q", core[16:24])[0]
    records = max(1, -(-total // size))
    if len(body) != total + 16 * records:
        raise ValueError("capsule is truncated or padded: header and body disagree")
    key = aes_key_unwrap(kek, wrapped)
    out = b""
    for index in range(records):
        record, body = body[:size + 16], body[size + 16:]
        out += AESGCM(key).decrypt(chunk_nonce(index, index == records - 1), record, core)
    return out


# --------------------------------------------------------------------------
# AB-R1 -- the photographic record and its Reed-Solomon protection
# --------------------------------------------------------------------------

RECORD_LEN, PARITY_LEN = 88, 40
CODEWORD_LEN = RECORD_LEN + PARITY_LEN            # 128 bytes = 1,024 embedded bits

_EXP, _LOG = [0] * 512, [0] * 256
_v = 1
for _i in range(255):
    _EXP[_i], _LOG[_v] = _v, _i
    _v <<= 1
    if _v & 0x100:
        _v ^= 0x11D                                # x^8 + x^4 + x^3 + x^2 + 1
for _i in range(255, 512):
    _EXP[_i] = _EXP[_i - 255]


def gf_mul(a, b):
    return 0 if a == 0 or b == 0 else _EXP[_LOG[a] + _LOG[b]]


def gf_div(a, b):
    return 0 if a == 0 else _EXP[(_LOG[a] - _LOG[b]) % 255]


def gf_pow(a, e):
    return _EXP[(_LOG[a] * e) % 255]


def poly_mul(p, q):
    r = [0] * (len(p) + len(q) - 1)
    for i, a in enumerate(p):
        for j, b in enumerate(q):
            r[i + j] ^= gf_mul(a, b)
    return r


def poly_eval(p, x):
    y = p[0]
    for c in p[1:]:
        y = gf_mul(y, x) ^ c
    return y


def _generator(nsym):
    g = [1]
    for i in range(nsym):
        g = poly_mul(g, [1, gf_pow(2, i)])         # (x - alpha^i), alpha = 2
    return g


_GEN = _generator(PARITY_LEN)


def rs_encode(msg: bytes) -> bytes:
    """Systematic RS over GF(2^8): codeword = message || 40 parity bytes."""
    buf = list(msg) + [0] * PARITY_LEN
    for i in range(len(msg)):
        c = buf[i]
        if c:
            for j in range(1, len(_GEN)):
                buf[i + j] ^= gf_mul(_GEN[j], c)
    return bytes(msg) + bytes(buf[len(msg):])


def rs_syndromes(cw):
    return [poly_eval(list(cw), gf_pow(2, i)) for i in range(PARITY_LEN)]


def rs_decode(cw: bytes) -> bytes:
    """Correct up to 20 byte errors. Berlekamp-Massey, Chien search, Forney."""
    cw = list(cw)
    synd = rs_syndromes(cw)
    if max(synd) == 0:
        return bytes(cw[:RECORD_LEN])
    # error locator sigma(x), lowest degree first
    sigma, prev, L, m, b = [1], [1], 0, 1, 1
    for r in range(PARITY_LEN):
        delta = synd[r]
        for i in range(1, L + 1):
            delta ^= gf_mul(sigma[i], synd[r - i])
        if delta == 0:
            m += 1
            continue
        scale = gf_div(delta, b)
        cand = sigma + [0] * max(0, len(prev) + m - len(sigma))
        for i, c in enumerate(prev):
            cand[i + m] ^= gf_mul(scale, c)
        if 2 * L <= r:
            prev, L, b, m = sigma, r + 1 - L, delta, 1
        else:
            m += 1
        sigma = cand
    sigma = sigma[:L + 1]
    n = len(cw)
    # Chien search: position j (from the left) has locator X = alpha^(n-1-j)
    places = [j for j in range(n) if poly_eval(sigma[::-1], gf_pow(2, (255 - (n - 1 - j)) % 255)) == 0]
    if len(places) != L:
        raise ValueError("record unrecoverable: more than 20 byte errors")
    # Forney, first consecutive root alpha^0
    omega = [0] * PARITY_LEN
    for i in range(PARITY_LEN):
        for j in range(min(i, L) + 1):
            omega[i] ^= gf_mul(sigma[j], synd[i - j])
    deriv = [sigma[i] if i % 2 == 1 else 0 for i in range(1, L + 1)]   # formal derivative
    for j in places:
        X = gf_pow(2, n - 1 - j)
        Xinv = gf_pow(2, (255 - (n - 1 - j)) % 255)
        num = poly_eval(omega[::-1], Xinv)
        den = poly_eval(deriv[::-1], Xinv) if deriv else 0
        if den == 0:
            raise ValueError("record unrecoverable")
        cw[j] ^= gf_mul(X, gf_div(num, den))
    if max(rs_syndromes(cw)) != 0:
        raise ValueError("record unrecoverable")
    return bytes(cw[:RECORD_LEN])


def record_build(photo, capsule, shard, wrapped_key, shard_sha256, flags) -> bytes:
    body = (b"AB8R" + bytes([1, flags]) + struct.pack(">H", photo) + bytes([capsule, shard, 0, 0])
            + wrapped_key + shard_sha256)
    assert len(body) == 84
    return body + struct.pack(">I", zlib.crc32(body))


def record_parse(record: bytes) -> dict:
    if len(record) != RECORD_LEN or record[:4] != b"AB8R":
        raise ValueError("not an AB-R1 record")
    if struct.unpack(">I", record[84:])[0] != zlib.crc32(record[:84]):
        raise ValueError("record CRC mismatch")
    flags = record[5]
    return {
        "profile": "AB-R1 v%d" % record[4],
        "wrapped_key_present": bool(flags & 1),
        "rehearsal": bool(flags & 2),
        "photograph": struct.unpack(">H", record[6:8])[0],
        "capsule": record[8],
        "shard": record[9],
        "wrapped_key": record[12:52].hex(),
        "shard_sha256": record[52:84].hex(),
    }


# --------------------------------------------------------------------------
# AB-S1 -- 88 shards, any 60
# --------------------------------------------------------------------------

SHARDS_K, SHARDS_N = 60, 88
# SHA-256 over the 88 shards of SHAKE256("afterberry/AB-S1/vector", 7777 bytes), record size 80
AB_S1_VECTOR = "ee0df4b6b0e99c05ef38cf8a07af4ba8393898b1b83c6d91766ba14e4dcea324"


def _mat_mul(a, b):
    return [[_dot(row, col) for col in zip(*b)] for row in a]


def _dot(row, col):
    acc = 0
    for u, v in zip(row, col):
        acc ^= gf_mul(u, v)
    return acc


def _mat_inv(m):
    """Gauss-Jordan over GF(2^8)."""
    size = len(m)
    work = [list(row) + [1 if i == j else 0 for j in range(size)] for i, row in enumerate(m)]
    for col in range(size):
        pivot = next(r for r in range(col, size) if work[r][col])
        work[col], work[pivot] = work[pivot], work[col]
        inv = gf_div(1, work[col][col])
        work[col] = [gf_mul(v, inv) for v in work[col]]
        for r in range(size):
            if r != col and work[r][col]:
                f = work[r][col]
                work[r] = [v ^ gf_mul(f, w) for v, w in zip(work[r], work[col])]
    return [row[size:] for row in work]


def shard_matrix():
    """88 x 60: Vandermonde V[r][c] = r^c, normalised so the top 60 rows are the identity."""
    vand = [[1 if c == 0 else (0 if r == 0 else gf_pow(r, c)) for c in range(SHARDS_K)]
            for r in range(SHARDS_N)]
    return _mat_mul(vand, _mat_inv(vand[:SHARDS_K]))


def shards_encode(ciphertext: bytes, record_size: int = CHUNK + 16):
    """Shards 1-60 are consecutive slices of the capsule ciphertext -- the chunk records,
    without the header core or the sealed key -- and every slice is a whole number of
    records, so a lone data shard can still be decrypted. Shards 61-88 are parity."""
    records = -(-len(ciphertext) // record_size)
    length = -(-records // SHARDS_K) * record_size
    padded = ciphertext + bytes(length * SHARDS_K - len(ciphertext))
    data = [padded[i * length:(i + 1) * length] for i in range(SHARDS_K)]
    rows = shard_matrix()[SHARDS_K:]
    parity = [bytes(_dot(row, column) for column in zip(*data)) for row in rows]
    return data + parity


def shard_file(core: bytes, wrapped: bytes, number: int, payload: bytes, ciphertext_len: int) -> bytes:
    """136-byte header, then the payload. Any one shard names its vault, capsule and lock."""
    head = (b"AB8S" + bytes([1, 0]) + core[6:8] + bytes([number, SHARDS_K, SHARDS_N, 0])
            + struct.pack(">QQ", len(payload), ciphertext_len) + core + wrapped)
    assert len(head) == 132
    return head + struct.pack(">I", zlib.crc32(head)) + payload


def shard_parse(data: bytes) -> dict:
    head = data[:132]
    if head[:4] != b"AB8S" or struct.unpack(">I", data[132:136])[0] != zlib.crc32(head):
        raise ValueError("not an AB-S1 shard, or header damaged")
    length, total = struct.unpack(">QQ", head[12:28])
    if head[4] != 1:
        raise ValueError("unknown shard profile version %d" % head[4])
    if (head[9], head[10]) != (head[28 + 8], head[28 + 9]) or head[6:8] != head[28 + 6:28 + 8]:
        raise ValueError("shard header disagrees with the capsule header core inside it")
    if len(data) != 136 + length:
        raise ValueError("shard is truncated or padded")
    return {"capsule": struct.unpack(">H", head[6:8])[0], "shard": head[8], "k": head[9], "n": head[10],
            "ciphertext_len": total, "core": head[28:92], "wrapped_key": head[92:132],
            "payload": data[136:136 + length]}


def shards_reconstruct(have: dict, total_length: int) -> bytes:
    """have maps shard number (1-88) to shard bytes; any 60 will do."""
    if len(have) < SHARDS_K:
        raise ValueError("need 60 shards, have %d" % len(have))
    numbers = sorted(have)[:SHARDS_K]
    matrix = shard_matrix()
    decode = _mat_inv([matrix[i - 1] for i in numbers])
    columns = list(zip(*[have[i] for i in numbers]))
    data = [bytes(_dot(row, column) for column in columns) for row in decode]
    return b"".join(data)[:total_length]


# --------------------------------------------------------------------------
# PNG, just enough: 8-bit RGB, non-interlaced
# --------------------------------------------------------------------------

_PNG_SIG = b"\x89PNG\r\n\x1a\n"


def png_read(path):
    data = open(path, "rb").read()
    if data[:8] != _PNG_SIG:
        raise ValueError("not a PNG file")
    pos, idat, width, height, text = 8, b"", None, None, {}
    while pos < len(data):
        length, kind = struct.unpack(">I4s", data[pos:pos + 8])
        body = data[pos + 8:pos + 8 + length]
        if zlib.crc32(kind + body) != struct.unpack(">I", data[pos + 8 + length:pos + 12 + length])[0]:
            raise ValueError("PNG chunk CRC mismatch")
        pos += 12 + length
        if kind == b"IHDR":
            width, height, depth, colour, _, _, interlace = struct.unpack(">IIBBBBB", body)
            if (depth, colour, interlace) != (8, 2, 0):
                raise ValueError("carrier must be 8-bit RGB, non-interlaced")
        elif kind == b"IDAT":
            idat += body
        elif kind == b"iTXt":
            key, _, rest = body.partition(b"\x00")
            text[key.decode("latin-1")] = rest[2:].split(b"\x00", 2)[2].decode("utf-8")
    raw = zlib.decompress(idat)
    stride, bpp = width * 3, 3
    pixels, prior = bytearray(), bytearray(stride)
    for y in range(height):
        f = raw[y * (stride + 1)]
        line = bytearray(raw[y * (stride + 1) + 1:(y + 1) * (stride + 1)])
        for i in range(stride):
            a = line[i - bpp] if i >= bpp else 0
            b = prior[i]
            c = prior[i - bpp] if i >= bpp else 0
            if f == 1:
                line[i] = (line[i] + a) & 255
            elif f == 2:
                line[i] = (line[i] + b) & 255
            elif f == 3:
                line[i] = (line[i] + (a + b) // 2) & 255
            elif f == 4:
                pa, pb, pc = abs(b - c), abs(a - c), abs(a + b - 2 * c)
                line[i] = (line[i] + (a if pa <= pb and pa <= pc else b if pb <= pc else c)) & 255
            elif f != 0:
                raise ValueError("unknown PNG filter")
        pixels += line
        prior = line
    return width, height, pixels, text


def png_write(path, width, height, pixels, record_hex=None):
    def chunk(kind, body):
        return struct.pack(">I", len(body)) + kind + body + struct.pack(">I", zlib.crc32(kind + body))
    stride = width * 3
    raw = b"".join(b"\x00" + bytes(pixels[y * stride:(y + 1) * stride]) for y in range(height))
    out = _PNG_SIG + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
    out += chunk(b"sRGB", b"\x00")
    if record_hex:
        out += chunk(b"iTXt", b"afterberry\x00\x00\x00\x00\x00" + record_hex.encode("ascii"))
    out += chunk(b"IDAT", zlib.compress(raw, 9)) + chunk(b"IEND", b"")
    open(path, "wb").write(out)


def pixel_digest(width, height, pixels) -> str:
    return hashlib.sha256(struct.pack(">II", width, height) + bytes(pixels)).hexdigest()


# --------------------------------------------------------------------------
# AB-C1 -- the carrier
# --------------------------------------------------------------------------

class _Stream:
    """SHA-256 in counter mode, read as big-endian 64-bit integers."""

    def __init__(self, seed: bytes):
        self.seed, self.counter, self.buf = seed, 0, []

    def next(self) -> int:
        if not self.buf:
            block = sha256(self.seed + struct.pack(">Q", self.counter))
            self.counter += 1
            self.buf = list(struct.unpack(">4Q", block))
        return self.buf.pop(0)


def _carrier_stream(photo: int, width: int, height: int) -> _Stream:
    return _Stream(sha256(b"afterberry/AB-C1/" + struct.pack(">HII", photo, width, height)))


def _positions(stream: _Stream, samples: int, count: int):
    if samples < 8 * count:
        raise ValueError("image too small to be a carrier")
    chosen, seen = [], set()
    while len(chosen) < count:
        s = stream.next() % samples
        if s not in seen:
            seen.add(s)
            chosen.append(s)
    return chosen


def _bits(data: bytes):
    return [(byte >> (7 - i)) & 1 for byte in data for i in range(8)]


def carrier_embed(pixels: bytearray, width: int, height: int, photo: int, codeword: bytes) -> int:
    stream = _carrier_stream(photo, width, height)
    changed = 0
    for s, bit in zip(_positions(stream, len(pixels), 8 * CODEWORD_LEN), _bits(codeword)):
        if pixels[s] & 1 == bit:
            continue
        step = 1 if stream.next() & 1 else -1
        if pixels[s] == 0:
            step = 1
        elif pixels[s] == 255:
            step = -1
        pixels[s] += step
        changed += 1
    return changed


def carrier_extract(pixels, width: int, height: int, photo: int) -> bytes:
    stream = _carrier_stream(photo, width, height)
    bits = [pixels[s] & 1 for s in _positions(stream, len(pixels), 8 * CODEWORD_LEN)]
    return bytes(sum(bit << (7 - i) for i, bit in enumerate(bits[j:j + 8])) for j in range(0, len(bits), 8))


# --------------------------------------------------------------------------
# command line
# --------------------------------------------------------------------------

def _cmd_lock_verify(args):
    doc = json.load(open(args.lockfile))
    lock = Lock(doc)
    print("lock        %s  (%d-bit modulus, %d waypoints)" % (lock.lock_id, lock.bits, len(lock.waypoints)))
    print("lock digest %s" % lock.digest.hex())
    extra = doc.get("rehearsal_only", {})
    if args.door == "capability":
        if "p_hex" not in extra:
            sys.exit("no factors in this lock file -- that is the point of the real one")
        walk = door_capability(lock, int(extra["p_hex"], 16), int(extra["q_hex"], 16))
    elif args.door == "patience":
        total = lock.waypoints[-1][1]
        print("squaring %s times; this is the slow door" % format(total, ","))
        started = time.time()
        walk = door_patience(lock, lambda d: print("  %5.1f%%  %6.0fs" % (100 * d / total, time.time() - started), end="\r"))
    else:
        a, done, ok = lock.x, 0, True
        for (ring, t, check), y_hex, proof in zip(lock.waypoints, extra["y_hex"], extra["relay_proofs"]):
            b = int(y_hex, 16)
            good = relay_verify(lock, a, b, t - done, int(proof["pi_hex"], 16)) and lock.check(ring, b) == check
            print("ring %d  leg of %s squarings  relay proof %s" % (ring, format(t - done, ","), "verifies" if good else "FAILS"))
            ok, a, done = ok and good, b, t
        sys.exit(0 if ok else 1)
    capsules = {c["ring"]: c for c in doc.get("capsules", [])}
    try:
        walk = list(walk) if args.door == "capability" else walk
        _report_rings(lock, walk, capsules, args.door == "patience")
    except ValueError as err:
        sys.exit("\nstopped: %s" % err)


def _report_rings(lock, walk, capsules, walked):
    for ring, t, y in walk:
        arrived = lock.check(ring, y) == dict((r, c) for r, _, c in lock.waypoints)[ring]
        stones = "; %d milestones verified on the way" % (t // lock.milestone_interval) if walked and lock.milestone_interval else ""
        print("ring %d  t = %s  waypoint %s%s" % (ring, format(t, ","), "reached" if arrived else "MISMATCH", stones))
        if arrived and ring in capsules:
            c = capsules[ring]
            kek = lock.capsule_kek(lock.ring_key(ring, y), c["capsule_id"])
            try:
                print("        “%s”" % open_capsule(bytes.fromhex(c["capsule_hex"]), kek).decode("utf-8"))
            except RuntimeError as err:
                print("        (%s)" % err)


def _cmd_carrier_extract(args):
    width, height, pixels, text = png_read(args.png)
    print("carrier      %d x %d, pixel digest %s" % (width, height, pixel_digest(width, height, pixels)))
    codeword = carrier_extract(pixels, width, height, args.index)
    errors = sum(1 for s in rs_syndromes(codeword) if s)
    record = rs_decode(codeword)
    print("embedded     %d bits read, %s" % (8 * CODEWORD_LEN, "clean" if not errors else "damaged, corrected"))
    if "afterberry" in text:
        print("iTXt copy    %s" % ("matches" if bytes.fromhex(text["afterberry"])[:RECORD_LEN] == record else "DIFFERS"))
    print(json.dumps(record_parse(record), indent=2))


def _cmd_carrier_embed(args):
    width, height, pixels, _ = png_read(args.source)
    codeword = rs_encode(bytes.fromhex(args.record))
    changed = carrier_embed(pixels, width, height, args.index, codeword)
    png_write(args.out, width, height, pixels, codeword.hex())
    if carrier_extract(pixels, width, height, args.index) != codeword:
        sys.exit("round trip failed")
    print("embedded 1,024 bits; %d samples changed by one level; round trip verified" % changed)


def _cmd_record_parse(args):
    data = bytes.fromhex(args.hex)
    record = rs_decode(data) if len(data) == CODEWORD_LEN else data
    print(json.dumps(record_parse(record), indent=2))


def _cmd_selftest(args):
    import os
    import random
    rng = random.Random(8888)
    # AB-S1: a known-answer vector, then lose any 28 of 88 shards, five times over
    vector = hashlib.shake_256(b"afterberry/AB-S1/vector").digest(7777)
    digest = hashlib.sha256(b"".join(shards_encode(vector, record_size=80))).hexdigest()
    assert digest == AB_S1_VECTOR, digest
    print("AB-S1  known-answer vector                                        ok")
    blob = os.urandom(7777)
    shards = shards_encode(blob, record_size=80)
    assert all(len(sh) % 80 == 0 for sh in shards)
    assert b"".join(shards[:SHARDS_K])[:len(blob)] == blob
    core = header_core(1, 101, SHARDS_K, SHARDS_N, 0, bytes(32))
    files = [shard_file(core, bytes(40), i + 1, sh, len(blob)) for i, sh in enumerate(shards)]
    parsed = shard_parse(files[87])
    assert (parsed["capsule"], parsed["shard"], parsed["payload"]) == (101, 88, shards[87])
    for _ in range(5):
        kept = rng.sample(range(1, SHARDS_N + 1), SHARDS_K)
        assert shards_reconstruct({i: shards[i - 1] for i in kept}, len(blob)) == blob
    print("AB-S1  88 shards, 28 lost at random, capsule rebuilt              ok")
    for bad in (files[0][:-1], files[0][:4] + bytes([9]) + files[0][5:]):
        try:
            shard_parse(bad)
            raise AssertionError("a damaged shard was accepted")
        except ValueError:
            pass
    print("AB-S1  shard header round trip; damaged shards refused            ok")
    # AB-E1: multi-chunk capsules, the final flag, truncation, reordering, the Remainder form
    try:
        kek, key = os.urandom(32), os.urandom(32)
        for size in (0, 1, 63, 64, 65, 200):
            plain = os.urandom(size)
            core = header_core(1, 7, SHARDS_K, SHARDS_N, size, bytes(32), chunk=64)
            capsule = seal_capsule(plain, key, kek, core)
            assert open_capsule(capsule, kek) == plain
        tampered = [capsule[:-80], capsule[:104] + capsule[184:264] + capsule[104:184] + capsule[264:]]
        for bad in tampered:                       # a record dropped; two records swapped
            try:
                open_capsule(bad, kek)
                raise AssertionError("a tampered capsule was accepted")
            except Exception as err:
                assert not isinstance(err, AssertionError)
        assert seal_capsule(plain, key, None, core)[64:104] == bytes(40)
        print("AB-E1  multi-chunk capsules; dropped and swapped records refused  ok")
    except RuntimeError as err:
        print("AB-E1  skipped: %s" % err)
    # AB-R1: twenty damaged bytes in a record codeword
    record = record_build(8888, 101, 88, os.urandom(40), os.urandom(32), 1)
    for _ in range(50):
        damaged = bytearray(rs_encode(record))
        for pos in rng.sample(range(CODEWORD_LEN), 20):
            damaged[pos] ^= rng.randrange(1, 256)
        assert rs_decode(bytes(damaged)) == record
    print("AB-R1  record codeword, 20 of 128 bytes damaged, corrected        ok")
    # AB-C1: embed and extract in a synthetic image, including saturated samples
    width, height = 64, 48
    pixels = bytearray(rng.choice((0, 255, rng.randrange(256))) for _ in range(width * height * 3))
    before = bytes(pixels)
    carrier_embed(pixels, width, height, 8888, rs_encode(record))
    assert rs_decode(carrier_extract(pixels, width, height, 8888)) == record
    assert max(abs(a - b) for a, b in zip(before, pixels)) == 1
    print("AB-C1  carrier round trip, no sample moved more than one level    ok")
    # AB-L1: a toy lock, both doors and a relay proof
    p, q = 1019 * 2 + 1, 1031 * 2 + 1          # 2039 and 2063: safe primes, far too small
    doc = {"lock_id": "selftest", "bits": (p * q).bit_length(), "n_hex": "%x" % (p * q),
           "waypoints": [{"ring": 1, "t": "5000", "check_sha256": "00" * 32}]}
    lock = Lock(doc)
    slow = list(door_patience(lock))[0][2]
    fast = list(door_capability(lock, p, q))[0][2]
    assert slow == fast
    assert not relay_verify(lock, lock.x, 0, 5000, 0)          # the degenerate proof is refused
    # (relay proofs need a modulus large enough for a 256-bit challenge to be meaningful;
    #  they are exercised against the rehearsal lock by: lock-verify --door proofs)
    print("AB-L1  patient door and capability door agree                     ok")


def main():
    ap = argparse.ArgumentParser(description="Afterberry reference tool, manual revision 2.0")
    sub = ap.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("lock-verify", help="open a lock file by one of its doors")
    s.add_argument("lockfile")
    s.add_argument("--door", choices=["capability", "patience", "proofs"], default="capability")
    s.set_defaults(fn=_cmd_lock_verify)
    s = sub.add_parser("carrier-extract", help="read the record out of a carrier PNG")
    s.add_argument("png")
    s.add_argument("--index", type=int, required=True, help="photograph index (0 = rehearsal)")
    s.set_defaults(fn=_cmd_carrier_extract)
    s = sub.add_parser("carrier-embed", help="write a record into an 8-bit RGB PNG")
    s.add_argument("source")
    s.add_argument("out")
    s.add_argument("--index", type=int, required=True)
    s.add_argument("--record", required=True, help="88-byte record, hex")
    s.set_defaults(fn=_cmd_carrier_embed)
    s = sub.add_parser("record-parse", help="decode an 88-byte record or 128-byte codeword")
    s.add_argument("hex")
    s.set_defaults(fn=_cmd_record_parse)
    s = sub.add_parser("selftest", help="exercise every profile with throwaway data")
    s.set_defaults(fn=_cmd_selftest)
    args = ap.parse_args()
    args.fn(args)


if __name__ == "__main__":
    main()