//
// Copyright (C) 2024 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use hex;
use kmr_common::crypto::{ConstTimeEq, MonotonicClock, Rng, Sha256};
use log::info;

/// Test basic hardware-backed key derivation functionality.
fn test_derive_hbk() {
    let key1 = crate::crypto::derive_hbk("Secret".as_bytes(), 16).unwrap();
    let key2 = crate::crypto::derive_hbk("Secret".as_bytes(), 16).unwrap();
    let key3 = crate::crypto::derive_hbk("Sekret".as_bytes(), 16).unwrap();
    assert_eq!(key1.0, key2.0);
    assert_ne!(key1.0, key3.0);
    info!("test_derive_hbk: Success");
}

/// Test basic [`MonotonicClock`] functionality.
pub fn test_clock<C: MonotonicClock>(clock: C) {
    let t1 = clock.now();
    assert!(t1.0 > 0);
    let t2 = clock.now();
    assert!(t2.0 >= t1.0);
    let t3 = clock.now();
    assert!(t3.0 >= t2.0);
    info!("test_clock: Success");
}

/// Test basic [`Rng`] functionality.
pub fn test_rng<R: Rng>(rng: &mut R) {
    let u1 = rng.next_u64();
    let u2 = rng.next_u64();
    assert_ne!(u1, u2);
    let mut b1 = [0u8; 16];
    let mut b2 = [0u8; 16];
    rng.fill_bytes(&mut b1);
    rng.fill_bytes(&mut b2);
    assert_ne!(b1, b2);
    rng.add_entropy(&b1);
    rng.add_entropy(&[]);
    rng.fill_bytes(&mut b1);
    assert_ne!(b1, b2);
    info!("test_rng: Success");
}

/// Test basic [`ConstTimeEq`] functionality. Does not test the key constant-time property though.
pub fn test_eq<E: ConstTimeEq>(comparator: E) {
    let b0 = [];
    let b1 = [0u8, 1u8, 2u8];
    let b2 = [1u8, 1u8, 2u8];
    let b3 = [0u8, 1u8, 3u8];
    let b4 = [0u8, 1u8, 2u8, 3u8];
    let b5 = [42; 4096];
    let mut b6 = [42; 4096];
    b6[4095] = 43;
    assert!(comparator.eq(&b0, &b0));
    assert!(comparator.eq(&b5, &b5));
    assert!(comparator.ne(&b0, &b1));
    assert!(comparator.ne(&b0, &b2));
    assert!(comparator.ne(&b0, &b3));
    assert!(comparator.ne(&b0, &b4));
    assert!(comparator.ne(&b0, &b5));
    assert!(comparator.eq(&b1, &b1));
    assert!(comparator.ne(&b1, &b2));
    assert!(comparator.ne(&b1, &b3));
    assert!(comparator.ne(&b5, &b4));
    assert!(comparator.ne(&b5, &b6));
    info!("test_eq: Success");
}

/// Test basic SHA-256 functionality.
pub fn test_sha256<S: Sha256>(sha256: S) {
    struct TestCase {
        msg: &'static [u8],
        want: &'static str,
    }
    let tests = vec![
        TestCase {
            msg: b"",
            want: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
        },
        TestCase {
            msg: b"abc",
            want: "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
        },
    ];
    for test in tests {
        let got = sha256.hash(test.msg).unwrap();
        assert_eq!(hex::encode(&got), test.want, "for input {}", hex::encode(test.msg));
    }
    info!("test_sha256: Success");
}

pub fn run() {
    info!("crypto::tests::run: Starting");

    test_derive_hbk();

    test_clock(crate::crypto::clock::Clock);

    test_eq(crate::crypto::eq::Eq);

    let mut rng = crate::crypto::rng::Rng;
    test_rng(&mut rng);

    test_sha256(crate::crypto::sha::Sha256 {});

    info!("crypto::tests::run: Done");
}