Files
tea3-py/src/tea3/tea3model.py
T
2026-04-21 14:44:35 +02:00

160 lines
5.2 KiB
Python

from sage.all import GF, BooleanPolynomialRing
from tea3.constants import TEA3_SBOX, T_F1, T_F2
from tea3.pretty_print import pretty_print, pretty_print_vec
class Tea3Model:
def __init__(self, max_steps=20):
self.F = GF(2)
self.step_count = 0
names = (
[f"x{i}{j}" for i in range(5) for j in range(8)] +
[f"r{i}{j}" for i in range(5) for j in range(8)] +
[f"R{i}{j}" for i in range(8) for j in range(8)] +
[f"f{s}_{i}{j}" for s in range(max_steps) for i in range(8) for j in range(8)]
)
name_string = ",".join(names)
self.S = BooleanPolynomialRing(len(names), name_string)
self.v = self.S.gens()
self.x_bits = [list(self.v[i*8:(i+1)*8]) for i in range(5)]
self.r_bits = [list(self.v[40 + i*8 : 40 + (i+1)*8]) for i in range(5)]
self.R_bits = [list(self.v[80 + i*8 : 80 + (i+1)*8]) for i in range(8)]
# Abstract variables
base = 80 + 64
self.fR_bits = [
[list(self.v[base + s*64 + i*8 : base + s*64 + i*8 + 8])
for i in range(8)]
for s in range(max_steps)
]
def _split_poly(self, poly):
"""
Split a polynomial into:
- R_f_part: monomials involving only 'R' or 'f' (abstract) variables
- xr_part: monomials involving 'x' or 'r' variables
constant term is grouped with R_f_part when R_f_part is non-zero
"""
zero = self.S.zero()
R_f_part = zero
xr_part = zero
has_const = bool(poly.constant_coefficient())
for monom in poly:
vars_in_term = monom.variables()
if not vars_in_term:
continue
families = {str(v)[0] for v in vars_in_term}
monom_poly = self.S(monom)
if families <= {'R', 'f'}:
R_f_part += monom_poly
else:
xr_part += monom_poly
if has_const:
if R_f_part != zero:
R_f_part += self.S.one()
else:
xr_part += self.S.one()
return R_f_part, xr_part
def _abstract_R(self):
"""
Replace the R/f-dependent part of every R_bits[i][j] with an
abstract variable f{step}_{i}{j}, leaving only x and r terms explicit.
"""
s = self.step_count
for i in range(8):
for j in range(8):
R_f_part, xr_part = self._split_poly(self.R_bits[i][j])
if R_f_part != self.S.zero():
self.R_bits[i][j] = self.fR_bits[s][i][j] + xr_part
else:
self.R_bits[i][j] = xr_part
def step(self):
R = self.R_bits.copy()
x = self.x_bits.copy()
r = self.r_bits.copy()
R0, R1, R2, R3, R4, R5, R6, R7 = R
x0, x1, x2, x3, x4 = x
r0, r1, r2, r3, r4 = r
# update r register
self.r_bits[0] = r1
self.r_bits[1] = r2
self.r_bits[2] = r3
self.r_bits[3] = r4
self.r_bits[4] = xor_vec(r0, S(r2))
# update x register
self.x_bits[0] = x1
self.x_bits[1] = x2
self.x_bits[2] = x3
self.x_bits[3] = x4
self.x_bits[4] = xor_vec(x0, r0)
# update R registers
self.R_bits[7] = R6
self.R_bits[6] = R5
self.R_bits[5] = xor_vec(R4, F31(R6, R5))
self.R_bits[4] = R3
self.R_bits[3] = R2
self.R_bits[2] = R1
self.R_bits[1] = R0
self.R_bits[0] = xor_vec(x0, xor_vec(R7, xor_vec(BP(R4), F32(R2, R1))))
self._abstract_R()
self.step_count += 1
return R7
def S(r):
# placeholder
return r
def xor_vec(a, b):
return [ai + bi for ai, bi in zip(a, b)]
def BP(r):
return [r[3], r[7], r[6], r[1], r[2], r[4], r[0], r[5]]
def F31(X, Y):
x0, x1, x2, x3, x4, x5, x6, x7 = X
y0, y1, y2, y3, y4, y5, y6, y7 = Y
return [
x5*x6*y5 + x5*x6 + x5*y5*y6 + x5*y5 + x5*y6 + x6*y5*y6 + y5*y6 + y5 + y6,
x6*x7*y6 + x6*x7 + x6*y6*y7 + x6*y7 + x6 + x7*y6 + x7*y7 + y6*y7 + y6 + y7 + 1,
x0*x7*y0 + x0*y0*y7 + x0*y0 + x0 + x7*y7 + y0,
x0*x1*y0 + x0*x1*y1 + x0*x1 + x0*y0*y1 + x0*y1 + x1*y0*y1 + x1 + y1,
x1*x2*y2 + x1*x2 + x1*y1 + x1*y2 + x2*y1*y2 + x2*y1 + x2 + y1*y2 + y1 + y2 + 1,
x2*x3*y3 + x2*y2 + x2 + x3*y2*y3 + x3*y3 + x3 + y2 + y3 + 1,
x3*x4*y4 + x3*y3 + x3*y4 + x4*y3*y4 + x4*y4 + x4 + y3*y4 + 1,
x4*x5*y5 + x4*y4*y5 + x4*y4 + x4*y5 + x5*y4*y5 + x5 + y5 + 1,
]
def F32(X, Y):
x0, x1, x2, x3, x4, x5, x6, x7 = X
y0, y1, y2, y3, y4, y5, y6, y7 = Y
return [
x5*x6*y5 + x5*x6 + x5*y5*y6 + x5*y5 + x5*y6 + x6*y5*y6 + y5 + y6 + 1,
x6*x7*y6 + x6*x7 + x6*y6*y7 + x6*y7 + x6 + x7*y7 + x7 + y6 + 1,
x0*x7*y0 + x0*x7*y7 + x0*y0*y7 + x0*y0 + x0 + x7*y0*y7 + y0,
x0*x1*y0 + x0*x1*y1 + x0*y0*y1 + x1*y0*y1 + x1*y0 + x1 + y0*y1 + y1,
x1*x2*y1 + x1*x2*y2 + x1*x2 + x1*y1*y2 + x1*y1 + x1 + x2*y1*y2 + x2 + y1*y2 + y2,
x2*x3*y3 + x2*y2 + x3*y2*y3 + x3*y3 + x3 + y3,
x3*x4*y4 + x3*x4 + x3*y3 + x3*y4 + x4*y3*y4 + x4*y3 + x4 + y3*y4 + y3 + y4,
x4*x5*y5 + x4*y4*y5 + x4*y4 + x5*y4*y5 + x5 + y4 + y5 + 1,
]