modelization + pretty print

src/tea3/pretty_print.py

@@ -0,0 +1,64 @@
+### print
+
+def monomial_to_str(exp, names):
+    factors = []
+    for i, e in enumerate(exp):
+        if e == 0:
+            continue
+        name = names[i]
+        if e == 1:
+            factors.append(name)
+        else:
+            factors.append(f"{name}^{e}")
+    return "*".join(factors) if factors else "1"
+
+
+def pretty_print(poly):
+    ring = poly.parent()
+    names = ring.variable_names()
+
+    R_terms = []
+    r_terms = []
+    x_terms = []
+    mixed_terms = []
+    has_const = False
+
+    for exp, coeff in poly.dict().items():
+        if coeff == 0:
+            continue
+
+        if sum(exp) == 0:
+            has_const = True
+            continue
+
+        term = monomial_to_str(exp, names)
+        vars_in_term = [names[i] for i, e in enumerate(exp) if e != 0]
+        families = {v[0] for v in vars_in_term}
+
+        if families == {"R"}:
+            R_terms.append(term)
+        elif families == {"r"}:
+            r_terms.append(term)
+        elif families == {"x"}:
+            x_terms.append(term)
+        else:
+            mixed_terms.append(term)
+
+    parts = []
+
+    if R_terms or (has_const and not r_terms):
+        parts.append("f(Ri)")
+    elif has_const:
+        parts.append("1")
+
+    if r_terms:
+        parts.append("f(ri)")
+
+    parts.extend(x_terms)
+    parts.extend(mixed_terms)
+
+    return " + ".join(parts) if parts else "0"
+
+
+def pretty_print_vec(vec):
+    return "[" + ",\n ".join(pretty_print(p) for p in vec) + "]"

src/tea3/tea3model.py

@@ -1,20 +1,105 @@
-from sage.all import *
+from sage.all import GF, PolynomialRing
 
 from tea3.constants import TEA3_SBOX, T_F1, T_F2
+from tea3.pretty_print import pretty_print, pretty_print_vec
 
 
-F = GF(2)
+class Tea3Model:
+    def __init__(self):
+        self.F = GF(2)
 
-names = (
+        names = (
-    [f"x{i}{j}" for i in range(5) for j in range(8)] +
+            [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(5) for j in range(8)] +
-    [f"R{i}{j}" for i in range(8) for j in range(8)]
+            [f"R{i}{j}" for i in range(8) for j in range(8)]
-)
+        )
 
-S = PolynomialRing(F, names)
+        self.S = PolynomialRing(self.F, names)
-v = S.gens()
+        self.v = self.S.gens()
 
-x_bits = [list(v[i*8:(i+1)*8]) for i in range(5)]
+        self.x_bits = [list(self.v[i*8:(i+1)*8]) for i in range(5)]
-r_bits = [list(v[40 + i*8 : 40 + (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)]
-R_bits = [list(v[80 + i*8 : 80 + (i+1)*8]) for i in range(8)]
+        self.R_bits = [list(self.v[80 + i*8 : 80 + (i+1)*8]) for i in range(8)]
 
+    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, G31(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), G32(R2, R1))))
+
+        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[2], r[7], r[3], r[5], r[6], r[1], r[0], r[4]]
+
+def G31(X, Y):
+    x3,x4,x5,x6,x7,x0,x1,x2 = X
+    y3,y4,y5,y6,y7,y0,y1,y2 = Y
+
+    return [
+        x1 + y0 + x0*y0 + y1 + x0*y1 + x0*x1*y1 + y0*y1 + x0*y0*y1 + x1*y0*y1,
+        1 + y1 + y2 + x1*y2 + x2*y2 + y1*y2 + x1*y1*y2,
+        x2 + x3 + x2*y2 + x3*y2 + y3 + x3*y3 + x3*y2*y3,
+        x4*y3 + y4 + x3*y3*y4 + x4*y3*y4,
+        1 + y4 + x4*y4 + x5*y4 + y5 + y4*y5 + x5*y4*y5,
+        1 + x5 + x6 + y5 + x5*y5 + x6*y5 + y6 + x6*y6 + x6*y5*y6,
+        1 + x6 + x6*y6 + x7*y7 + y6*y7 + x7*y6*y7,
+        1 + x7 + y0 + x0*y7 + x7*y7 + x0*x7*y7 + x0*y0*y7 + x7*y0*y7,
+    ]
+
+def G32(X, Y):
+    x3,x4,x5,x6,x7,x0,x1,x2 = X
+    y3,y4,y5,y6,y7,y0,y1,y2 = Y
+
+    return [
+        1 + x1 + x0*x1 + y0 + x0*y0 + x1*y0 + y1 + x0*x1*y1 + x0*y0*y1 + x1*y0*y1,
+        1 + y1 + x2*y2 + x1*y1*y2,
+        x3 + x3*y2 + y3 + x2*y3 + x3*y3 + x2*y2*y3 + x3*y2*y3,
+        x4*y3 + y4 + y3*y4 + x3*y3*y4 + x4*y3*y4,
+        x4*y4 + y5 + y4*y5 + x4*y4*y5 + x5*y4*y5,
+        x5 + x6 + x5*y5 + x6*y5 + y6 + x6*y6 + x6*y5*y6,
+        y6 + x6*y6 + x7*y6 + y7 + y6*y7 + x7*y6*y7,
+        x0*x7 + y0 + x7*y0 + y7 + x0*y7 + x7*y7 + x0*x7*y7 + x0*y0*y7 + x7*y0*y7,
+    ]
+
+
+
+t = Tea3Model()
+for i in range(3):
+    print("step "+str(i))
+    t.step()
+    print(pretty_print_vec(t.R_bits[0]))