//! This example implements the PQC KAT KEM tests.
//! It generates `.req` and `.rsp` files which can be compared
//! to the files generated by the C reference implementation.

use rusty_saber::api::{
    CRYPTO_ALGNAME, CRYPTO_BYTES, CRYPTO_CIPHERTEXTBYTES, CRYPTO_PUBLICKEYBYTES,
    CRYPTO_SECRETKEYBYTES,
};
use rusty_saber::kem::{crypto_kem_dec, crypto_kem_enc, crypto_kem_keypair};
use rusty_saber::rng::{AesState, RNGState};
use std::error::Error;
use std::fs::File;
use std::io::{prelude::*, BufRead, BufReader, SeekFrom};
use std::str;

fn main() -> Result<(), Box<dyn Error>> {
    let mut seed = vec![0u8; 48];
    let mut entropy_input = [0u8; 48];
    let (mut ct, mut ss, mut ss1) = (
        [0u8; CRYPTO_CIPHERTEXTBYTES],
        [0u8; CRYPTO_BYTES],
        [0u8; CRYPTO_BYTES],
    );
    let mut count_buf: Vec<u8>;
    let (mut pk, mut sk) = ([0u8; CRYPTO_PUBLICKEYBYTES], [0u8; CRYPTO_SECRETKEYBYTES]);

    //Create the REQUEST file
    let fn_req = format!("PQCkemKAT_{}.req", CRYPTO_SECRETKEYBYTES);
    let mut req_file = File::create(fn_req.to_string())?;

    //Create the RESPONSE file
    let fn_rsp = format!("PQCkemKAT_{}.rsp", CRYPTO_SECRETKEYBYTES);
    let mut rsp_file = File::create(fn_rsp)?;

    let mut rng_state = AesState::with_increasing_seed();

    //Fill the .req file
    for i in 0..100 {
        write!(req_file, "count = {}\n", i)?;
        rng_state.randombytes(&mut seed)?;
        fprint_bstr(&mut req_file, "seed = ".to_string(), &seed, 48)?;
        write!(req_file, "pk =\n")?;
        write!(req_file, "sk =\n")?;
        write!(req_file, "ct =\n")?;
        write!(req_file, "ss =\n\n")?;
    }
    let req_file_read = File::open(fn_req.to_string())?;

    //Fill the .rsp file
    write!(rsp_file, "# {}\n\n", CRYPTO_ALGNAME)?;

    let mut req_reader = BufReader::new(req_file_read);

    loop {
        count_buf = vec![];
        if find_marker(&mut req_reader, "count = ")? {
            req_reader.read_until(b'\n', &mut count_buf)?;
        } else {
            break;
        }
        write!(rsp_file, "count = {}", str::from_utf8(&*count_buf)?)?;

        read_hex(&mut req_reader, &mut seed, 48, "seed = ")?;
        fprint_bstr(&mut rsp_file, "seed = ".to_string(), &seed, 48)?;

        let mut rng = AesState::new();
        entropy_input.copy_from_slice(&seed[..]);
        rng.randombytes_init(entropy_input);

        // Generate the public/private keypair
        crypto_kem_keypair(&mut pk, &mut sk, &mut rng)?;

        fprint_bstr(
            &mut rsp_file,
            "pk = ".to_string(),
            &pk,
            CRYPTO_PUBLICKEYBYTES,
        )?;
        fprint_bstr(
            &mut rsp_file,
            "sk = ".to_string(),
            &sk,
            CRYPTO_SECRETKEYBYTES,
        )?;

        crypto_kem_enc(&mut ct, &mut ss, &mut pk, &mut rng)?;

        fprint_bstr(
            &mut rsp_file,
            "ct = ".to_string(),
            &ct,
            CRYPTO_CIPHERTEXTBYTES,
        )?;
        fprint_bstr(&mut rsp_file, "ss = ".to_string(), &ss, CRYPTO_BYTES)?;
        write!(rsp_file, "\n")?;

        crypto_kem_dec(&mut ss1, &ct, &sk)?;
        assert_eq!(ss1, ss, "crypto_kem_dec returned bad 'ss' value\n");
    }
    Ok(())
}

fn find_marker(reader: &mut BufReader<File>, marker: &str) -> Result<bool, Box<dyn Error>> {
    let mut line = String::new();
    while let Ok(num) = reader.read_line(&mut line) {
        if num == 0 {
            break;
        }
        if line.starts_with(marker) {
            reader.seek(SeekFrom::Current(-((line.len() - marker.len()) as i64)))?;
            return Ok(true);
        }
        line.clear();
    }

    Ok(false)
}

fn read_hex(
    reader: &mut BufReader<File>,
    a: &mut Vec<u8>,
    length: usize,
    str: &str,
) -> Result<bool, Box<dyn Error>> {
    let mut tmp_ch = vec![0u8; 1];
    let mut ch: u8;
    let mut started: bool = false;
    let mut ich: char;
    if length == 0 {
        a[0] = 0x00;
        return Ok(true);
    }
    a.fill(0);

    if find_marker(reader, str)? {
        loop {
            match reader.read_exact(&mut tmp_ch) {
                Ok(_) => {}
                Err(_) => {
                    break;
                }
            }
            ch = tmp_ch[0];
            if !ch.is_ascii_hexdigit() {
                if !started {
                    if ch == '\n' as u8 {
                        break;
                    } else {
                        continue;
                    }
                } else {
                    break;
                }
            }

            started = true;

            if (ch >= '0' as u8) && (ch <= '9' as u8) {
                ich = (ch - '0' as u8) as char;
            } else if (ch >= 'A' as u8) && (ch <= 'F' as u8) {
                ich = (ch - 'A' as u8 + 10) as char;
            } else if (ch >= 'a' as u8) && (ch <= 'f' as u8) {
                ich = (ch - 'a' as u8 + 10) as char;
            } else {
                ich = 0 as char;
            }

            for i in 0..length - 1 {
                a[i] = (a[i] << 4) | (a[i + 1] >> 4);
            }
            a[length - 1] = (a[length - 1] << 4) | ich as u8;
        }
    } else {
        return Ok(false);
    }

    Ok(true)
}

fn fprint_bstr(file: &mut File, s: String, a: &[u8], l: usize) -> Result<(), Box<dyn Error>> {
    write!(file, "{}", s)?;

    for i in 0..l {
        write!(file, "{:02X}", a[i])?;
    }

    if l == 0 {
        write!(file, "00")?;
    }

    write!(file, "\n")?;
    Ok(())
}