"""AES-128 encryption of one block, written from FIPS 197, and counter mode from NIST SP 800-38A
§6.5, in pure Python, for the prior-carriage verifier's sealed sets (prior-carriage §18.2).

It shares no code with wai-rs, which uses a Rust AES implementation. Only the forward cipher is
needed: counter mode encrypts and decrypts with it alike.
"""


def _xtime(a):
    a <<= 1
    return (a ^ 0x11B) & 0xFF if a & 0x100 else a


def _mul(a, b):
    r = 0
    while b:
        if b & 1:
            r ^= a
        a = _xtime(a)
        b >>= 1
    return r


def _sbox():
    """FIPS 197 §5.1.1: the multiplicative inverse in GF(2^8), then the affine transformation."""
    inv = [0] * 256
    for a in range(1, 256):
        for b in range(1, 256):
            if _mul(a, b) == 1:
                inv[a] = b
                break
    box = []
    for a in range(256):
        b = inv[a]
        s = b
        for k in range(1, 5):
            s ^= ((b << k) | (b >> (8 - k))) & 0xFF
        box.append(s ^ 0x63)
    return box


SBOX = _sbox()


def expand_key(key):
    """FIPS 197 §5.2 for a 16-byte key: 11 round keys of 16 bytes."""
    assert len(key) == 16
    w = [list(key[4 * i:4 * i + 4]) for i in range(4)]
    rcon = 1
    for i in range(4, 44):
        t = list(w[i - 1])
        if i % 4 == 0:
            t = [SBOX[x] for x in t[1:] + t[:1]]
            t[0] ^= rcon
            rcon = _xtime(rcon)
        w.append([a ^ b for a, b in zip(w[i - 4], t)])
    return [sum(w[4 * r:4 * r + 4], []) for r in range(11)]


def encrypt_block(round_keys, block):
    """FIPS 197 §5.1: the state in column order, ten rounds."""
    s = [b ^ k for b, k in zip(block, round_keys[0])]
    for r in range(1, 11):
        s = [SBOX[x] for x in s]
        # ShiftRows: byte (row, column) at index 4 * column + row moves left by its row.
        s = [s[4 * ((c + row) % 4) + row] for c in range(4) for row in range(4)]
        if r != 10:
            out = []
            for c in range(4):
                a = s[4 * c:4 * c + 4]
                out += [_mul(a[0], 2) ^ _mul(a[1], 3) ^ a[2] ^ a[3],
                        a[0] ^ _mul(a[1], 2) ^ _mul(a[2], 3) ^ a[3],
                        a[0] ^ a[1] ^ _mul(a[2], 2) ^ _mul(a[3], 3),
                        _mul(a[0], 3) ^ a[1] ^ a[2] ^ _mul(a[3], 2)]
            s = out
        s = [b ^ k for b, k in zip(s, round_keys[r])]
    return bytes(s)


def ctr(key, iv, data):
    """Counter mode: the counter block from `iv`, incremented as one 128-bit big-endian integer."""
    rk = expand_key(key)
    counter = int.from_bytes(iv, "big")
    out = bytearray()
    for i in range(0, len(data), 16):
        ks = encrypt_block(rk, counter.to_bytes(16, "big"))
        out += bytes(a ^ b for a, b in zip(data[i:i + 16], ks))
        counter = (counter + 1) % (1 << 128)
    return bytes(out)


def self_test():
    """FIPS 197 Appendix C.1 and NIST SP 800-38A F.5.1's first two blocks."""
    rk = expand_key(bytes.fromhex("000102030405060708090a0b0c0d0e0f"))
    if encrypt_block(rk, bytes.fromhex("00112233445566778899aabbccddeeff")).hex() != "69c4e0d86a7b0430d8cdb78070b4c55a":
        return False
    key = bytes.fromhex("2b7e151628aed2a6abf7158809cf4f3c")
    iv = bytes.fromhex("f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff")
    plain = bytes.fromhex("6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e51")
    return ctr(key, iv, plain).hex() == "874d6191b620e3261bef6864990db6ce9806f66b7970fdff8617187bb9fffdff"