turing-cipher 0.1.1

use super::turing_tables::{MULTAB, QBOX, SBOX};

const MAX_KEY_LENGTH: usize = 32; // bytes
const MAX_IV_LENGTH: usize = 48; // bytes
const MAX_STREAM_LENGTH: usize = 340; // bytes, maximum stream generated by one call
const SHIFT_REGISTER_LENGTH: usize = 17; // words
const ROUND_OUTPUT_LENGTH: usize = 20; // bytes

pub struct Turing {
    key_length: usize,
    mixed_key: [u32; MAX_KEY_LENGTH / 4],
    shift_register: [u32; SHIFT_REGISTER_LENGTH],
    s0: [u32; 256],
    s1: [u32; 256],
    s2: [u32; 256],
    s3: [u32; 256],
}

impl Turing {
    pub fn new() -> Self {
        Turing {
            key_length: 0,
            mixed_key: [0; MAX_KEY_LENGTH / 4],
            shift_register: [0; SHIFT_REGISTER_LENGTH],
            s0: [0; 256],
            s1: [0; 256],
            s2: [0; 256],
            s3: [0; 256],
        }
    }

    /// 设置密钥
    pub fn turing_key(&mut self, key: &[u8], length: usize) -> Result<(), &'static str> {
        if (length & 0x03) != 0 || length > MAX_KEY_LENGTH {
            return Err("Invalid key length");
        }

        self.key_length = 0;
        for i in (0..length).step_by(4) {
            self.mixed_key[self.key_length] = Self::fixed_s(Self::byte_array_to_word(key, i));
            self.key_length += 1;
        }
        Self::mix_words(&mut self.mixed_key, self.key_length);

        self.build_s_box_tables();
        Ok(())
    }

    /// 设置初始化向量
    pub fn turing_iv(&mut self, iv: &[u8], length: usize) -> Result<(), &'static str> {
        if (length & 0x03) != 0 || (length + 4 * self.key_length) > MAX_IV_LENGTH {
            return Err("Invalid IV length");
        }

        let mut i = 0;
        let mut j = 0;

        // 首先复制IV并进行混合
        while j < length {
            self.shift_register[i] = Self::fixed_s(Self::byte_array_to_word(iv, j));
            i += 1;
            j += 4;
        }

        // 继续使用预混合的密钥
        for j in 0..self.key_length {
            self.shift_register[i] = self.mixed_key[j];
            i += 1;
        }

        // 长度相关的字
        self.shift_register[i] =
            ((self.key_length as u32) << 4) | ((length as u32) >> 2) | 0x01020300;
        i += 1;

        // 填充寄存器的其余部分
        let mut j = 0;
        while i < SHIFT_REGISTER_LENGTH {
            self.shift_register[i] = self.s(
                self.shift_register[j].wrapping_add(self.shift_register[i - 1]),
                0,
            );
            i += 1;
            j += 1;
        }

        // 混合所有字
        Self::mix_words(&mut self.shift_register, SHIFT_REGISTER_LENGTH);
        Ok(())
    }

    /// Generate a 340-byte buffer of cipher data.
    /// Return the number of bytes generated
    pub fn turing_gen(&mut self, buf: &mut [u8]) -> Result<usize, &'static str> {
        if buf.len() < MAX_STREAM_LENGTH {
            return Err("Buffer is too small");
        }

        let rounds = [
            (0, 0), (5, 20), (10, 40),
            (15, 60), (3, 80), (8, 100),
            (13, 120), (1, 140), (6, 160),
            (11, 180), (16, 200), (4, 220),
            (9, 240), (14, 260), (2, 280),
            (7, 300), (12, 320),
        ];

        for (z, offset) in rounds {
            self.turing_gen_round(z, buf, offset);
        }

        Ok(MAX_STREAM_LENGTH)
    }

    /// 字节数组转字(大端序)
    fn byte_array_to_word(b: &[u8], offset: usize) -> u32 {
        ((b[offset] as u32) << 24)
            | ((b[offset + 1] as u32) << 16)
            | ((b[offset + 2] as u32) << 8)
            | (b[offset + 3] as u32)
    }

    fn build_s_box_tables(&mut self) {
        for j in 0..256 {
            let mut w = 0;
            let mut k = j as u32;
            for i in 0..self.key_length {
                k = SBOX[(Self::get_byte(self.mixed_key[i], 0) ^ k) as usize];
                w ^= Self::left_rotate_word(QBOX[k as usize], i as u32);
            }
            self.s0[j] = (w & 0x00FFFFFF) | (k << 24);
        }

        for j in 0..256 {
            let mut w = 0;
            let mut k = j as u32;
            for i in 0..self.key_length {
                k = SBOX[(Self::get_byte(self.mixed_key[i], 1) ^ k) as usize];
                w ^= Self::left_rotate_word(QBOX[k as usize], (i + 8) as u32);
            }
            self.s1[j] = (w & 0xFF00FFFF) | (k << 16);
        }

        for j in 0..256 {
            let mut w = 0;
            let mut k = j as u32;
            for i in 0..self.key_length {
                k = SBOX[(Self::get_byte(self.mixed_key[i], 2) ^ k) as usize];
                w ^= Self::left_rotate_word(QBOX[k as usize], (i + 16) as u32);
            }
            self.s2[j] = (w & 0xFFFF00FF) | (k << 8);
        }

        for j in 0..256 {
            let mut w = 0;
            let mut k = j as u32;
            for i in 0..self.key_length {
                k = SBOX[(Self::get_byte(self.mixed_key[i], 3) ^ k) as usize];
                w ^= Self::left_rotate_word(QBOX[k as usize], (i + 24) as u32);
            }
            self.s3[j] = (w & 0xFFFFFF00) | k;
        }
    }

    /// 获取字的指定字节
    fn get_byte(word: u32, i: u32) -> u32 {
        (word >> (24 - 8 * i)) & 0xff
    }

    /// 字左循环移位
    fn left_rotate_word(word: u32, bits: u32) -> u32 {
        word.rotate_left(bits)
    }

    /// Performs a reversible transformation of a word, based on the S-boxes.
    /// The reversibility isn't used, but it guarantees no loss of information,
    /// and hence no equivalent keys or IVs.
    fn fixed_s(w: u32) -> u32 {
        let mut w = w;
        let mut b;

        b = SBOX[Self::get_byte(w, 0) as usize];
        w = ((w ^ QBOX[b as usize]) & 0x00ffffff) | (b << 24);

        b = SBOX[Self::get_byte(w, 1) as usize];
        w = ((w ^ Self::left_rotate_word(QBOX[b as usize], 8)) & 0xff00ffff) | (b << 16);

        b = SBOX[Self::get_byte(w, 2) as usize];
        w = ((w ^ Self::left_rotate_word(QBOX[b as usize], 16)) & 0xffff00ff) | (b << 8);

        b = SBOX[Self::get_byte(w, 3) as usize];
        w = ((w ^ Self::left_rotate_word(QBOX[b as usize], 24)) & 0xffffff00) | b;

        w
    }

    /// Pushes a word through the keyed S-boxes.
    /// As the bytes bounce around the permutation table, they are used
    /// to build up words from the Qbox entries. Then the byte position
    /// corresponding to the input byte is replaced with the result of
    /// the S-box, which is a permutation of the input and guarantees
    /// a balanced function.
    /// Also added a rotation of the input word, to combat a differential
    /// trail allowed by the PHT.
    fn s(&self, w: u32, r: u32) -> u32 {
        self.s0[Self::get_byte(w, (r) & 0x3) as usize]
            ^ self.s1[Self::get_byte(w, (1 + r) & 0x3) as usize]
            ^ self.s2[Self::get_byte(w, (2 + r) & 0x3) as usize]
            ^ self.s3[Self::get_byte(w, (3 + r) & 0x3) as usize]
    }

    /// General word-wide n-PHT (Pseudo-Hadamard Transform)
    fn mix_words(w: &mut [u32], n: usize) {
        let mut sum: u32 = 0;
        let mut i = 0;

        while i < n - 1 {
            sum = sum.wrapping_add(w[i]);
            i += 1;
        }
        w[n - 1] = w[n - 1].wrapping_add(sum);
        sum = w[n - 1];

        i = 0;
        while i < n - 1 {
            w[i] = w[i].wrapping_add(sum);
            i += 1;
        }
    }

    /// Generates a 5-word block of output.
    /// Buffering issues are outside the scope of this implementation.
    /// Returns the number of bytes of stream generated.
    fn turing_gen_round(&mut self, z: usize, buf: &mut [u8], offset: usize) -> usize {
        self.step(z);

        let mut a = self.shift_register[Self::offset(z + 1, 16)];
        let mut b = self.shift_register[Self::offset(z + 1, 13)];
        let mut c = self.shift_register[Self::offset(z + 1, 6)];
        let mut d = self.shift_register[Self::offset(z + 1, 1)];
        let mut e = self.shift_register[Self::offset(z + 1, 0)];

        // 混合5个字(变种PHT)
        e = e.wrapping_add(a.wrapping_add(b).wrapping_add(c).wrapping_add(d));
        a = a.wrapping_add(e);
        b = b.wrapping_add(e);
        c = c.wrapping_add(e);
        d = d.wrapping_add(e);

        a = self.s(a, 0);
        b = self.s(b, 1);
        c = self.s(c, 2);
        d = self.s(d, 3);
        e = self.s(e, 0);

        // 再次混合5个字
        e = e.wrapping_add(a.wrapping_add(b).wrapping_add(c).wrapping_add(d));
        a = a.wrapping_add(e);
        b = b.wrapping_add(e);
        c = c.wrapping_add(e);
        d = d.wrapping_add(e);

        self.step(z + 1);
        self.step(z + 2);
        self.step(z + 3);

        a = a.wrapping_add(self.shift_register[Self::offset(z + 4, 14)]);
        b = b.wrapping_add(self.shift_register[Self::offset(z + 4, 12)]);
        c = c.wrapping_add(self.shift_register[Self::offset(z + 4, 8)]);
        d = d.wrapping_add(self.shift_register[Self::offset(z + 4, 1)]);
        e = e.wrapping_add(self.shift_register[Self::offset(z + 4, 0)]);

        Self::word_to_byte_array(a, buf, offset);
        Self::word_to_byte_array(b, buf, offset + 4);
        Self::word_to_byte_array(c, buf, offset + 8);
        Self::word_to_byte_array(d, buf, offset + 12);
        Self::word_to_byte_array(e, buf, offset + 16);

        self.step(z + 4);

        ROUND_OUTPUT_LENGTH
    }

    /// 步进LFSR
    fn step(&mut self, z: usize) {
        let sr15 = self.shift_register[Self::offset(z, 15)];
        let sr4 = self.shift_register[Self::offset(z, 4)];
        let sr0 = self.shift_register[Self::offset(z, 0)];

        self.shift_register[Self::offset(z, 0)] =
            sr15 ^ sr4 ^ (sr0 << 8) ^ MULTAB[((sr0 >> 24) & 0xFF) as usize];
    }

    /// 计算移位寄存器的当前偏移量
    fn offset(zero: usize, i: usize) -> usize {
        (zero + i) % SHIFT_REGISTER_LENGTH
    }

    /// 将字转换为字节数组(大端序)
    fn word_to_byte_array(word: u32, b: &mut [u8], offset: usize) {
        b[offset] = (word >> 24) as u8;
        b[offset + 1] = (word >> 16) as u8;
        b[offset + 2] = (word >> 8) as u8;
        b[offset + 3] = word as u8;
    }
}