/******************************************************************************/
/* SlunkCrypt, by LoRd_MuldeR <MuldeR2@GMX.de>                                */
/* This work has been released under the CC0 1.0 Universal license!           */
/******************************************************************************/

use slunkcrypt_rs::{SlunkCryptPasswd, SlunkCrypt};
use hex;
use memx;
use rand::{rngs::OsRng, RngCore};
use std::{env, mem, cmp::Ordering};

const DEFAULT_TEST_LOOPS: usize = 8;
const PASSPHRASE: SlunkCryptPasswd     = SlunkCryptPasswd::Str("eUGAwxOm0QwiVnDx");
const PASSPHRASE_BAD: SlunkCryptPasswd = SlunkCryptPasswd::Str("eUGAwxOm1QwiVnDx");

#[test]
fn test_create_instance() {
    run_test(|| {
        let (context, nonce) = SlunkCrypt::init_encrypt(&PASSPHRASE, None).expect("Failed to create encryption context!");
        println!("EncryptContext: {:?}", context);
        let context = SlunkCrypt::init_decrypt(&PASSPHRASE, nonce, None).expect("Failed to create decryption context!");
        println!("DecryptContext: {:?}", context);
    });
}

#[test]
fn test_process() {
    run_test(|| {
        let (mut orig, mut encr, mut decr) = ([ 0u8; 5003 ], [ 0u8; 5003 ], [ 0u8; 5003 ]);
        OsRng.fill_bytes(&mut orig);
        println!("Plaintext: {}", hex::encode(&orig));

        let (mut context_enc, nonce) = SlunkCrypt::init_encrypt(&PASSPHRASE, None).expect("Failed to create encryption context!");
        context_enc.process(&orig, &mut encr).expect("Failed to encrypt data!");
        mem::drop(context_enc);

        println!("Encrypted: {}", hex::encode(&encr));
        assert!(memx::memcmp(&orig, &encr) != Ordering::Equal);

        let mut context_dec = SlunkCrypt::init_decrypt(&PASSPHRASE, nonce, None).expect("Failed to create decryption context!");
        context_dec.process(&encr, &mut decr).expect("Failed to encrypt data!");
        mem::drop(context_dec);

        println!("Decrypted: {}", hex::encode(&decr));
    });
}

#[test]
fn test_inplace() {
    run_test(|| {
        let (mut data, mut orig) = ([ 0u8; 5003 ], [ 0u8; 5003 ]);
        OsRng.fill_bytes(&mut data);
        memx::memcpy(&mut orig, &data).expect("Failed to copy original data!");

        println!("Plaintext: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) == Ordering::Equal);

        let (mut context_enc, nonce) = SlunkCrypt::init_encrypt(&PASSPHRASE, None).expect("Failed to create encryption context!");
        context_enc.inplace(&mut data).expect("Failed to encrypt data!");
        mem::drop(context_enc);

        println!("Encrypted: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) != Ordering::Equal);

        let mut context_dec = SlunkCrypt::init_decrypt(&PASSPHRASE, nonce, None).expect("Failed to create decryption context!");
        context_dec.inplace(&mut data).expect("Failed to decrypt data!");
        mem::drop(context_dec);

        println!("Decrypted: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) == Ordering::Equal);
    });
}

#[test]
fn test_reset() {
    run_test(|| {
        let (mut data, mut orig) = ([ 0u8; 5003 ], [ 0u8; 5003 ]);
        OsRng.fill_bytes(&mut data);
        memx::memcpy(&mut orig, &data).expect("Failed to copy original data!");

        println!("Plaintext: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) == Ordering::Equal);

        let (mut context, _nonce) = SlunkCrypt::init_encrypt(&PASSPHRASE, None).expect("Failed to create encryption context!");
        context.inplace(&mut data).expect("Failed to encrypt data!");

        println!("Encrypted: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) != Ordering::Equal);

        memx::memcpy(&mut data, &orig).expect("Failed to copy original data!");
        println!("Plaintext: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) == Ordering::Equal);

        let nonce = context.reset_encrypt(&PASSPHRASE).expect("Failed to reset encryption context!");
        context.inplace(&mut data).expect("Failed to encrypt data!");
    
        println!("Encrypted: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) != Ordering::Equal);

        context.reset_decrypt(&PASSPHRASE, nonce).expect("Failed to reset decryption context!");
        context.inplace(&mut data).expect("Failed to decrypt data!");
        mem::drop(context);

        println!("Decrypted: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) == Ordering::Equal);
    });
}

#[test]
fn test_fail_bad_nonce() {
    run_test(|| {
        let (mut data, mut orig) = ([ 0u8; 5003 ], [ 0u8; 5003 ]);
        OsRng.fill_bytes(&mut data);
        memx::memcpy(&mut orig, &data).expect("Failed to copy original data!");

        println!("Plaintext: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) == Ordering::Equal);
    
        let (mut context_enc, nonce) = SlunkCrypt::init_encrypt(&PASSPHRASE, None).expect("Failed to create encryption context!");
        context_enc.inplace(&mut data).expect("Failed to encrypt data!");
        mem::drop(context_enc);

        println!("Encrypted: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) != Ordering::Equal);

        let mut context_dec = SlunkCrypt::init_decrypt(&PASSPHRASE, nonce + 1, None).expect("Failed to create decryption context!");
        context_dec.inplace(&mut data).expect("Failed to decrypt data!");
        mem::drop(context_dec);

        println!("Decrypted: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) != Ordering::Equal);
    });
}

#[test]
fn test_fail_bad_passwd() {
    run_test(|| {
        let (mut data, mut orig) = ([ 0u8; 5003 ], [ 0u8; 5003 ]);
        OsRng.fill_bytes(&mut data);
        memx::memcpy(&mut orig, &data).expect("Failed to copy original data!");

        println!("Plaintext: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) == Ordering::Equal);

        let (mut context_enc, nonce) = SlunkCrypt::init_encrypt(&PASSPHRASE, None).expect("Failed to create encryption context!");
        context_enc.inplace(&mut data).expect("Failed to encrypt data!");
        mem::drop(context_enc);

        println!("Encrypted: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) != Ordering::Equal);

        let mut context_dec = SlunkCrypt::init_decrypt(&PASSPHRASE_BAD, nonce, None).expect("Failed to create decryption context!");
        context_dec.inplace(&mut data).expect("Failed to decrypt data!");
        mem::drop(context_dec);

        println!("Decrypted: {}", hex::encode(&data));
        assert!(memx::memcmp(&data, &orig) != Ordering::Equal);
    });
}

#[test]
fn test_stress() {
    let (mut data, mut orig) = ([ 0u8; u16::MAX as usize ], [ 0u8; u16::MAX  as usize ]);

    let (mut context_enc, nonce) = SlunkCrypt::init_encrypt(&PASSPHRASE, None).expect("Failed to create encryption context!");
    let mut context_dec = SlunkCrypt::init_decrypt(&PASSPHRASE, nonce,   None).expect("Failed to create decryption context!");
    for i in 0..u16::MAX {
        OsRng.fill_bytes(&mut data);
        memx::memcpy(&mut orig, &data).expect("Failed to copy original data!");

        context_enc.inplace(&mut data).expect("Failed to encrypt data!");
        assert!(memx::memcmp(&data, &orig) != Ordering::Equal);

        context_dec.inplace(&mut data).expect("Failed to decrypt data!");    
        assert!(memx::memcmp(&data, &orig) == Ordering::Equal);

        if i % 499 == 0 {
            println!("{:.1}%", (i as f64 / u16::MAX as f64) * 100.0)
        }
    }
}

fn run_test<F>(func: F) where F: Fn() {
    let loops = env::var("TEST_LOOPS").map_or(DEFAULT_TEST_LOOPS, |value| value.parse::<usize>().unwrap());
    for _ in 0..loops {
         func();
    }
}