def expand_key(key):
"""Key expansion routine to generate a key schedule.
Expand a short key into a single array of round keys, each 4 words.
The key expansion generates a total of Nb * (Nr + 1) 4-byte words:
- 128 -> 4 * (10 + 1) -> 44
- 192 -> 4 * (12 + 1) -> 52
- 256 -> 4 * (14 + 1) -> 60
This diverges from the specification in that they expanded key is
*not* a parameter that can be modified and remember its state.
Rather, it gets created here as an empty array and then filled.
"""
# The first *nk* words of the expanded key are filled with
# respective keys from the cipher key. (Each word == 4 bytes)
nk = int(key.size / 4)
nr = numrounds(nk)
nwords = NB * (nr + 1)
w = np.empty((NB, nwords), dtype=uint8)
w[:, :nk] = key
# Now fill out the rest of w.
# The dreaded loop:
# It may be near-impossible to do this without a loop, since
# we're dependent on the value w[i - nk]
i = nk
while i < nwords:
# Starts on row 3
temp = w[:, i - 1]
if i % nk == 0:
temp = xor(
# XOR of a 4-byte word with length-4 lookup from RCON table
sub_word(rot_word(temp)),
RCON[int(i / nk)]
)
elif nk > 6 and i % nk == 4:
temp = sub_word(temp)
w[:, i] = xor(w[:, i - nk], temp)
i += 1
return w
def decrypt_raw(state, key):
"""Decrypt a single ciphertext data block, `state`, using `key`.
Modifies `state` in-place!
Parameters
----------
inp: np.ndarray
key: np.ndarray
Returns
-------
state: np.ndarray
"""
nk = int(key.size / 4)
nr = numrounds(nk)
exkeys = expand_key(key) # 4 rows, NB * (nr + 1) columns (words)
exkeys = np.split(exkeys, exkeys.shape[1] / NB, axis=1)
# First XOR is with just input + key (reverse order of roundkeys)
# Last round doesn't get an InvMixColumns
xor(state, exkeys[-1], out=state)
inv_shift_rows(state, out=state)
inv_sub_bytes(state, out=state)
for ek in exkeys[nr - 1:0:-1]:
xor(state, ek, out=state)
inv_mix_columns(state, out=state)
inv_shift_rows(state, out=state)
inv_sub_bytes(state, out=state)
# One final (inverse) xor
xor(state, exkeys[0], out=state)
return state
def encrypt_raw(state, key):
"""Encrypt a single input data block, `state`, using `key`.
Modifies `state` in-place!
Parameters
----------
inp: np.ndarray
key: np.ndarray
Returns
-------
state: np.ndarray
"""
# We get the keys in "human form" and apply AES' column-based axis swap on them
# (This should happen within `encrypt()`)
# Retain the 'original' key for our expand_key methodology
# key = key.swapaxes(0, 1)
nk = int(key.size / 4)
nr = numrounds(nk)
exkeys = expand_key(key) # 4 rows, NB * (nr + 1) columns (words)
exkeys = np.split(exkeys, exkeys.shape[1] / NB, axis=1)
# First XOR is with just input + key
xor(state, exkeys[0], out=state)
# Intermediate rounds
for ek in exkeys[1:nr]:
sub_bytes(state, out=state)
shift_rows(state, out=state)
mix_columns(state, out=state)
xor(state, ek, out=state)
# Final round before a final XOR. No mixColumns here
sub_bytes(state, out=state)
shift_rows(state, out=state)
xor(state, exkeys[-1], out=state)
# TODO: this is returned in column-ordered format.
# We will need to massage it back into row ordered before flattening
return state
df = pd.DataFrame(
product(*vars.values()),
columns = vars.keys()
)
df['Pr'] = df.apply(lambda args: prob(*args), axis=1)
return df
if __name__=='__main__':
# Model 1: p=1/2, q=1/2
# U ~ Ber(p) - True if allergic
# X ~ Ber(q) - treatment
# Y = U xor X - recovers
model1 = lambda u, x, y: [1-p, p][u] * [1-q, q][x] * int(y==xor(u,x))
# Model 2: p=1/2, q=1/4, r=3/4
# U ~ Ber(p) - True if allergic
# X ~ Ber(q if U else r) - treatment
# Y = U xor X - recovers
model2 = lambda u, x, y: [1-p, p][u] * [[1-r, 1-q], [r, q]][x][u] * int(y==xor(u,x))
# Model variables
vars_ = dict(zip(['U', 'X', 'Y'], [[False, True], [False, True], [False, True]]))
# Model parameters
p, q, r, = symbols('p q r', real=True, nonnegative=True)
# Compute
df1 = get_joint(model1, **vars_)
def test_rcon_lookup(asw_to_axo):
"""Test of intermediate values from xor with RCON lookup."""
for i, (asw, axo) in enumerate(asw_to_axo, 1):
assert array_equal(xor(asw, RCON[i]), axo), i