aes.py

# -*- coding: utf-8 -*-
"""AES-128 logic that makes its work visible.

Instead of encrypting silently, ``encrypt_steps`` returns a list of step
objects. Each step describes exactly one transformation (state before, state
after, changed cells, a short explanation). The GUI only iterates over this
list and computes nothing itself.

Shape of a step::

    Step = {
        "name":         "SubBytes",        # SubBytes | ShiftRows | MixColumns | AddRoundKey
        "round":        3,                 # 0..10
        "state_before": [[..4..], ...],    # 4x4 matrix before
        "state_after":  [[..4..], ...],    # 4x4 matrix after
        "round_key":    [[..4..], ...] | None,  # only set for AddRoundKey
        "changed":      [(row, col), ...], # changed cells (for highlighting)
        "explanation":  "short text",      # shown by the GUI
        "mix_terms":    terms | None,      # only set for MixColumns, see mix_columns_terms()
    }

Every value here is JSON-serialisable as-is, which is what lets the web UI
just ``json.dumps`` the whole step list and render it in the browser.
"""

from sbox import SBOX, RCON


# ---------------------------------------------------------------------------
# GF(2^8) arithmetic
# ---------------------------------------------------------------------------

def xtime(b):
    """Multiply a byte by 2 in GF(2^8), reduction polynomial 0x11B."""
    b <<= 1
    if b & 0x100:
        b ^= 0x11B
    return b & 0xFF


def gf_mul(a, b):
    """Multiply two bytes in GF(2^8) (Russian peasant method via xtime)."""
    result = 0
    for _ in range(8):
        if b & 1:
            result ^= a
        b >>= 1
        a = xtime(a)
    return result & 0xFF


# ---------------------------------------------------------------------------
# State helpers (a state is a 4x4 matrix of bytes)
# ---------------------------------------------------------------------------

def bytes_to_state(data):
    """Fill 16 bytes column by column into a 4x4 matrix.

    The rule is: state[row][col] = data[col * 4 + row].
    """
    state = [[0] * 4 for _ in range(4)]
    for col in range(4):
        for row in range(4):
            state[row][col] = data[col * 4 + row]
    return state


def state_to_bytes(state):
    """Read a 4x4 matrix back into 16 bytes, column by column."""
    data = []
    for col in range(4):
        for row in range(4):
            data.append(state[row][col])
    return data


def copy_state(state):
    """Return a deep copy of a 4x4 matrix."""
    return [list(row) for row in state]


def diff_cells(before, after):
    """Return all cells (row, col) that changed between before and after."""
    changed = []
    for row in range(4):
        for col in range(4):
            if before[row][col] != after[row][col]:
                changed.append((row, col))
    return changed


# ---------------------------------------------------------------------------
# Key expansion
# ---------------------------------------------------------------------------

def _rot_word(word):
    """RotWord: [a, b, c, d] -> [b, c, d, a]."""
    return word[1:] + word[:1]


def _sub_word(word):
    """SubWord: replace each byte of the word through the S-box."""
    return [SBOX[b] for b in word]


def key_expansion(key):
    """Expand the 16-byte key into 11 round keys (44 words).

    Returns ``(round_keys, words)``:
      * round_keys: list of 11 4x4 matrices
      * words:      list of 44 word vectors (4 bytes each)
    """
    words = []
    # W[0..3] are simply the original key.
    for i in range(4):
        words.append([key[4 * i + 0], key[4 * i + 1], key[4 * i + 2], key[4 * i + 3]])

    # W[4..43] are derived from their predecessors.
    for i in range(4, 44):
        temp = list(words[i - 1])
        if i % 4 == 0:
            temp = _sub_word(_rot_word(temp))
            rc = RCON[i // 4]
            temp = [temp[j] ^ rc[j] for j in range(4)]
        words.append([words[i - 4][j] ^ temp[j] for j in range(4)])

    # Four consecutive words form the columns of a round key (column-wise).
    round_keys = []
    for r in range(11):
        rk = [[0] * 4 for _ in range(4)]
        for col in range(4):
            word = words[r * 4 + col]
            for row in range(4):
                rk[row][col] = word[row]
        round_keys.append(rk)
    return round_keys, words


# ---------------------------------------------------------------------------
# The four transformations
# ---------------------------------------------------------------------------

def sub_bytes(state):
    """SubBytes: replace each byte through the S-box (out = SBOX[in])."""
    result = copy_state(state)
    for row in range(4):
        for col in range(4):
            result[row][col] = SBOX[state[row][col]]
    return result


def shift_rows(state):
    """ShiftRows: rotate row r left by r positions (row 0 stays)."""
    result = copy_state(state)
    for row in range(4):
        result[row] = state[row][row:] + state[row][:row]
    return result


# The fixed MixColumns matrix (each row dotted with a column gives one output byte).
MIX_MATRIX = [[2, 3, 1, 1], [1, 2, 3, 1], [1, 1, 2, 3], [3, 1, 1, 2]]


def mix_columns(state):
    """MixColumns: multiply each column by the fixed matrix in GF(2^8)."""
    result = [[0] * 4 for _ in range(4)]
    for col in range(4):
        s = [state[row][col] for row in range(4)]
        for row in range(4):
            value = 0
            for k in range(4):
                value ^= gf_mul(MIX_MATRIX[row][k], s[k])
            result[row][col] = value
    return result


def mix_columns_terms(state):
    """Return the GF(2^8) terms behind every MixColumns output byte.

    terms[row][col] is a list of 4 ``(coefficient, input_byte, product)``
    tuples -- XORing the four products yields ``result[row][col]``. Used by
    the UI to show the arithmetic instead of just the answer.
    """
    terms = [[None] * 4 for _ in range(4)]
    for col in range(4):
        s = [state[row][col] for row in range(4)]
        for row in range(4):
            terms[row][col] = [
                (MIX_MATRIX[row][k], s[k], gf_mul(MIX_MATRIX[row][k], s[k]))
                for k in range(4)
            ]
    return terms


def add_round_key(state, round_key):
    """AddRoundKey: XOR the state with the round key, byte by byte."""
    result = copy_state(state)
    for row in range(4):
        for col in range(4):
            result[row][col] = state[row][col] ^ round_key[row][col]
    return result


# ---------------------------------------------------------------------------
# Step creation and overall flow
# ---------------------------------------------------------------------------

def _make_step(name, round_index, before, after, round_key, explanation):
    """Build a step with deep copies and automatic diff marking."""
    return {
        "name": name,
        "round": round_index,
        "state_before": copy_state(before),
        "state_after": copy_state(after),
        "round_key": copy_state(round_key) if round_key is not None else None,
        "changed": diff_cells(before, after),
        "explanation": explanation,
        "mix_terms": mix_columns_terms(before) if name == "MixColumns" else None,
    }


def encrypt_steps(plaintext, key):
    """Encrypt a 16-byte block and return ``(steps, ciphertext)``.

    Arguments:
      * plaintext: 16 bytes (bytes or list of ints)
      * key:       16 bytes (bytes or list of ints)

    Returns:
      * steps:      list of step dicts (see module docstring)
      * ciphertext: 16 bytes as a list of ints

    AES-128 produces exactly 40 steps:
      1 (round 0: AddRoundKey) + 9 * 4 (rounds 1-9) + 3 (round 10).
    """
    plaintext = list(plaintext)
    key = list(key)
    if len(plaintext) != 16 or len(key) != 16:
        raise ValueError("Plaintext and key must each be 16 bytes long.")

    round_keys, _ = key_expansion(key)
    steps = []
    state = bytes_to_state(plaintext)

    # Round 0: just AddRoundKey with the first round key (the key itself).
    after = add_round_key(state, round_keys[0])
    steps.append(_make_step(
        "AddRoundKey", 0, state, after, round_keys[0],
        "Start: the plaintext is XORed with the first round key (the key itself).",
    ))
    state = after

    # Rounds 1-9: SubBytes -> ShiftRows -> MixColumns -> AddRoundKey.
    for r in range(1, 10):
        before = state
        after = sub_bytes(before)
        steps.append(_make_step(
            "SubBytes", r, before, after, None,
            "Each byte is replaced through the S-box (non-linear substitution).",
        ))
        state = after

        before = state
        after = shift_rows(before)
        steps.append(_make_step(
            "ShiftRows", r, before, after, None,
            "Row r is cyclically rotated left by r positions "
            "(row 0 stays, row 1 by 1, row 2 by 2, row 3 by 3).",
        ))
        state = after

        before = state
        after = mix_columns(before)
        steps.append(_make_step(
            "MixColumns", r, before, after, None,
            "Each column is multiplied in GF(2^8) by the fixed MixColumns matrix "
            "(mixing within the column).",
        ))
        state = after

        before = state
        after = add_round_key(before, round_keys[r])
        steps.append(_make_step(
            "AddRoundKey", r, before, after, round_keys[r],
            "The state is XORed with this round's round key.",
        ))
        state = after

    # Round 10 (final round): SubBytes -> ShiftRows -> AddRoundKey, no MixColumns!
    before = state
    after = sub_bytes(before)
    steps.append(_make_step(
        "SubBytes", 10, before, after, None,
        "Final round: each byte is substituted through the S-box one last time.",
    ))
    state = after

    before = state
    after = shift_rows(before)
    steps.append(_make_step(
        "ShiftRows", 10, before, after, None,
        "Final round: the rows are rotated one last time.",
    ))
    state = after

    before = state
    after = add_round_key(before, round_keys[10])
    steps.append(_make_step(
        "AddRoundKey", 10, before, after, round_keys[10],
        "Final round: AddRoundKey with no following MixColumns yields the "
        "finished ciphertext.",
    ))
    state = after

    ciphertext = state_to_bytes(state)
    return steps, ciphertext


# ---------------------------------------------------------------------------
# Self-test against the FIPS-197 test vector (Appendix B)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    plaintext = [0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d,
                 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37, 0x07, 0x34]
    key = [0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
           0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c]
    expected = [0x39, 0x25, 0x84, 0x1d, 0x02, 0xdc, 0x09, 0xfb,
                0xdc, 0x11, 0x85, 0x97, 0x19, 0x6a, 0x0b, 0x32]

    steps, ciphertext = encrypt_steps(plaintext, key)
    assert len(steps) == 40, f"Expected 40 steps, got {len(steps)}."
    assert ciphertext == expected, "Ciphertext does not match FIPS-197!"

    hex_ct = " ".join(f"{b:02x}" for b in ciphertext)
    print("OK - FIPS-197 Appendix B test vector passed.")
    print(f"OK - Ciphertext: {hex_ct}")
    print(f"OK - generated {len(steps)} steps.")