#![allow(dead_code)]

mod aes;
mod magic;

pub mod rsa_utils;

use aes::*;
use magic::*;
use rand_mt::Mt64;

pub struct MhyXorpad([u8; 4096]);

impl MhyXorpad {
    pub fn from_ec2b(ec2b: &[u8]) -> Self {
        let ec2b_magic = u32::from_be_bytes(ec2b[0..4].try_into().unwrap());
        assert_eq!(ec2b_magic, 0x45633262); // "Ec2b"

        let key_len = u32::from_le_bytes(ec2b[4..8].try_into().unwrap());
        assert_eq!(key_len, 16);

        let data_len = u32::from_le_bytes(ec2b[24..28].try_into().unwrap());
        assert_eq!(data_len, 2048);

        let mut key: [u8; 16] = ec2b[8..24].try_into().unwrap();
        let data = &ec2b[28..2048 + 28];

        key_scramble(&mut key);
        (0..16).for_each(|i| key[i] ^= KEY_XORPAD_TABLE[i]);

        Self(get_decrypt_vector(&mut key, data))
    }

    pub fn from_seed(seed: u64) -> Self {
        let mut mt = Mt64::new(seed);
        let new_seed = mt.next_u64();
        mt.reseed(new_seed);
        mt.next_u64();

        let mut buf = Vec::with_capacity(4096);
        (0..512).for_each(|_| buf.extend(mt.next_u64().to_be_bytes()));

        Self(buf.try_into().unwrap())
    }

    pub fn xor(&self, data: &mut [u8]) {
        for (i, b) in data.iter_mut().enumerate() {
            *b ^= self.0[i % 4096]
        }
    }
}

fn key_scramble(key: &mut [u8]) {
    let mut round_keys = [0u8; 176];
    for round in 0..11 {
        for i in 0..16 {
            for j in 0..16 {
                let idx = (round << 8) + (i * 16) + j;
                round_keys[round * 16 + i] ^= AES_XORPAD_TABLE[1][idx] ^ AES_XORPAD_TABLE[0][idx];
            }
        }
    }

    let mut chip = [0u8; 16];
    oqs_mhy128_enc_c(key, &round_keys, &mut chip);

    key.copy_from_slice(&chip);
}

fn get_decrypt_vector(key: &mut [u8], crypt: &[u8]) -> [u8; 4096] {
    let mut output = Vec::with_capacity(4096);

    let val = crypt
        .chunks_exact(8)
        .map(|chunk| u64::from_le_bytes(chunk.try_into().unwrap()))
        .fold(0xFFFFFFFFFFFFFFFF, |val, i| val ^ i);

    let key_qword_0 = u64::from_le_bytes(key[0..8].try_into().unwrap());
    let key_qword_1 = u64::from_le_bytes(key[8..16].try_into().unwrap());

    let mut mt = Mt64::new(key_qword_1 ^ 0xCEAC3B5A867837AC ^ val ^ key_qword_0);
    for _ in 0..512 {
        output.extend(mt.next_u64().to_le_bytes());
    }

    output.try_into().unwrap()
}
