107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Sequence
|
|
|
|
from tea3.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
|