initial commit

src/tea3/tea3.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Sequence
+
+from .constants import TEA3_LUT_A, TEA3_LUT_B, TEA3_SBOX
+
+MASK32 = 0xFFFFFFFF
+MASK64 = 0xFFFFFFFFFFFFFFFF
+
+
+def _ror8(x: int, n: int = 1) -> int:
+    x &= 0xFF
+    return ((x >> n) | ((x << (8 - n)) & 0xFF)) & 0xFF
+
+
+def tea3_reorder_state_byte(b_st_byte: int) -> int:
+    b_st_byte &= 0xFF
+    b_out = 0
+    b_out |= (b_st_byte << 6) & 0x40
+    b_out |= (b_st_byte << 1) & 0x20
+    b_out |= (b_st_byte << 2) & 0x98
+    b_out |= (b_st_byte >> 4) & 0x04
+    b_out |= (b_st_byte >> 3) & 0x01
+    b_out |= (b_st_byte >> 6) & 0x02
+    return b_out & 0xFF
+
+
+def tea3_state_word_to_newbyte(w_st: int, aw_lut: Sequence[int]) -> int:
+    w_st &= 0xFFFF
+    b_st0 = w_st & 0xFF
+    b_st1 = (w_st >> 8) & 0xFF
+    b_out = 0
+
+    for i in range(8):
+        b_dist = ((b_st0 >> 5) & 0x03) | ((b_st1 >> 3) & 0x0C)
+        if aw_lut[i] & (1 << b_dist):
+            b_out |= 1 << i
+        b_st0 = _ror8(b_st0, 1)
+        b_st1 = _ror8(b_st1, 1)
+
+    return b_out & 0xFF
+
+
+@dataclass
+class Tea3:
+    key_register: list[int]
+    iv: int
+    next_skip_rounds: int
+
+    def __init__(self, frame_number: int, key_register: Sequence[int]):
+        if len(key_register) != 10:
+            raise ValueError("TEA3 key register must be 10 bytes")
+        self.key_register = [b & 0xFF for b in key_register]
+        self.iv = self.compute_iv(frame_number)
+        self.next_skip_rounds = 51
+
+    @staticmethod
+    def compute_iv(frame_number: int) -> int:
+        frame_number &= MASK32
+        dw_xorred = (frame_number ^ 0xC43A7D51) & MASK32
+        dw_xorred = ((dw_xorred << 8) | (dw_xorred >> 24)) & MASK32
+        qw_iv = ((frame_number << 32) | dw_xorred) & MASK64
+        qw_iv = ((qw_iv >> 8) | ((qw_iv & 0xFF) << 56)) & MASK64
+        return qw_iv
+
+    def reset(self, frame_number: int, key_register: Sequence[int]) -> None:
+        if len(key_register) != 10:
+            raise ValueError("TEA3 key register must be 10 bytes")
+        self.key_register = [b & 0xFF for b in key_register]
+        self.iv = self.compute_iv(frame_number)
+        self.next_skip_rounds = 51
+
+    def step(self) -> None:
+        # Step 1: key schedule byte
+        b_sbox_out = TEA3_SBOX[self.key_register[7] ^ self.key_register[2]] ^ self.key_register[0]
+
+        self.key_register = self.key_register[1:] + [b_sbox_out & 0xFF]
+
+        # Step 2: state-derived bytes
+        b_deriv_byte12 = tea3_state_word_to_newbyte((self.iv >> 8) & 0xFFFF, TEA3_LUT_A)
+        b_deriv_byte56 = tea3_state_word_to_newbyte((self.iv >> 40) & 0xFFFF, TEA3_LUT_B)
+        b_reord_byte4 = tea3_reorder_state_byte((self.iv >> 32) & 0xFF)
+
+        # Step 3: combine
+        b_new_byte = (((self.iv >> 56) & 0xFF) ^ b_reord_byte4 ^ b_deriv_byte12 ^ b_sbox_out) & 0xFF
+        b_mix_byte = b_deriv_byte56 & 0xFF
+
+        # Step 4: update 64-bit state
+        self.iv = (((self.iv << 8) & MASK64) ^ ((b_mix_byte & 0xFF) << 40)) | b_new_byte
+        self.iv &= MASK64
+
+    def next_byte(self) -> int:
+        for _ in range(self.next_skip_rounds):
+            self.step()
+        self.next_skip_rounds = 19
+        return (self.iv >> 56) & 0xFF
+
+    def keystream(self, length: int) -> list[int]:
+        return [self.next_byte() for _ in range(length)]
+
+    def key_register_view(self) -> tuple[int, ...]:
+        return tuple(self.key_register)
+
+    def iv_view(self) -> int:
+        return self.iv