#![cfg(any(kdf_srtp, all(hmac, any(kdf_pbkdf2, kdf_ssh, kdf_tls13))))]

#[cfg(all(hmac, any(kdf_pbkdf2, kdf_tls13)))]
use wolfssl_wolfcrypt::hmac::HMAC;
use wolfssl_wolfcrypt::kdf::*;
#[cfg(all(hmac, kdf_tls13))]
use wolfssl_wolfcrypt::sha::SHA256;

#[test]
#[cfg(all(hmac, kdf_pbkdf2))]
fn test_pbkdf2() {
    let password = b"passwordpassword";
    let salt = [0x78u8, 0x57, 0x8E, 0x5a, 0x5d, 0x63, 0xcb, 0x06];
    let iterations = 2048;
    let expected_key = [
        0x43u8, 0x6d, 0xb5, 0xe8, 0xd0, 0xfb, 0x3f, 0x35, 0x42, 0x48, 0x39, 0xbc,
        0x2d, 0xd4, 0xf9, 0x37, 0xd4, 0x95, 0x16, 0xa7, 0x2a, 0x9a, 0x21, 0xd1
    ];

    let mut keyout = [0u8; 24];
    pbkdf2(password, &salt, iterations, HMAC::TYPE_SHA256, &mut keyout).expect("Error with pbkdf2()");
    assert_eq!(keyout, expected_key);

    let mut keyout = [0u8; 24];
    pbkdf2_ex(password, &salt, iterations, HMAC::TYPE_SHA256, None, None, &mut keyout).expect("Error with pbkdf2_ex()");
    assert_eq!(keyout, expected_key);
}

#[test]
#[cfg(all(hmac, kdf_pbkdf2))]
fn test_pkcs12_pbkdf() {
    let password = [0x00u8, 0x73, 0x00, 0x6d, 0x00, 0x65, 0x00, 0x67, 0x00, 0x00];
    let salt = [0x0au8, 0x58, 0xCF, 0x64, 0x53, 0x0d, 0x82, 0x3f];
    let expected_key = [
        0x27u8, 0xE9, 0x0D, 0x7E, 0xD5, 0xA1, 0xC4, 0x11,
        0xBA, 0x87, 0x8B, 0xC0, 0x90, 0xF5, 0xCE, 0xBE,
        0x5E, 0x9D, 0x5F, 0xE3, 0xD6, 0x2B, 0x73, 0xAA
    ];
    let iterations = 1;

    let mut keyout = [0u8; 24];
    pkcs12_pbkdf(&password, &salt, iterations, HMAC::TYPE_SHA256, 1, &mut keyout).expect("Error with pkcs12_pbkdf()");
    assert_eq!(keyout, expected_key);

    let mut keyout = [0u8; 24];
    pkcs12_pbkdf_ex(&password, &salt, iterations, HMAC::TYPE_SHA256, 1, None, &mut keyout).expect("Error with pkcs12_pbkdf_ex()");
    assert_eq!(keyout, expected_key);
}

#[test]
#[cfg(all(hmac, kdf_tls13))]
fn test_tls13_hkdf_extract_expand() {
    let hash_hello1 = [
        0x63u8, 0x83, 0x58, 0xab, 0x36, 0xcd, 0x0c, 0xf3,
        0x26, 0x07, 0xb5, 0x5f, 0x0b, 0x8b, 0x45, 0xd6,
        0x7d, 0x5b, 0x42, 0xdc, 0xa8, 0xaa, 0x06, 0xfb,
        0x20, 0xa5, 0xbb, 0x85, 0xdb, 0x54, 0xd8, 0x8b
    ];
    let client_early_traffic_secret = [
        0x20u8, 0x18, 0x72, 0x7c, 0xde, 0x3a, 0x85, 0x17, 0x72, 0xdc, 0xd7, 0x72,
        0xb0, 0xfc, 0x45, 0xd0, 0x62, 0xb9, 0xbb, 0x38, 0x69, 0x05, 0x7b, 0xb4,
        0x5e, 0x58, 0x5d, 0xed, 0xcd, 0x0b, 0x96, 0xd3
    ];

    let mut secret = [0u8; SHA256::DIGEST_SIZE];

    tls13_hkdf_extract(HMAC::TYPE_SHA256, None, None, &mut secret).expect("Error with tls13_hkdf_extract()");
    tls13_hkdf_extract_ex(HMAC::TYPE_SHA256, None, None, &mut secret, None, None).expect("Error with tls13_hkdf_extract_ex()");

    let protocol_label = b"tls13 ";
    let ce_traffic_label = b"c e traffic";
    let mut expand_out = [0u8; SHA256::DIGEST_SIZE];

    tls13_hkdf_expand_label(HMAC::TYPE_SHA256, &secret,
        protocol_label, ce_traffic_label,
        &hash_hello1, &mut expand_out).expect("Error with tls13_hkdf_expand_label()");
    tls13_hkdf_expand_label_ex(HMAC::TYPE_SHA256, &secret,
        protocol_label, ce_traffic_label,
        &hash_hello1, &mut expand_out, None, None).expect("Error with tls13_hkdf_expand_label_ex()");

    assert_eq!(expand_out, client_early_traffic_secret);
}

#[test]
#[cfg(all(hmac, kdf_ssh))]
fn test_ssh_kdf() {
    let ssh_kdf_set3_k = [
        0x6Au8, 0xC3, 0x82, 0xEA, 0xAC, 0xA0, 0x93, 0xE1,
        0x25, 0xE2, 0x5C, 0x24, 0xBE, 0xBC, 0x84, 0x64,
        0x0C, 0x11, 0x98, 0x75, 0x07, 0x34, 0x4B, 0x5C,
        0x73, 0x9C, 0xEB, 0x84, 0xA9, 0xE0, 0xB2, 0x22,
        0xB9, 0xA8, 0xB5, 0x1C, 0x83, 0x9E, 0x5E, 0xBE,
        0x49, 0xCF, 0xAD, 0xBF, 0xB3, 0x95, 0x99, 0x76,
        0x4E, 0xD5, 0x22, 0x09, 0x9D, 0xC9, 0x12, 0x75,
        0x19, 0x50, 0xDC, 0x7D, 0xC9, 0x7F, 0xBD, 0xC0,
        0x63, 0x28, 0xB6, 0x8F, 0x22, 0x78, 0x1F, 0xD3,
        0x15, 0xAF, 0x56, 0x80, 0x09, 0xA5, 0x50, 0x9E,
        0x5B, 0x87, 0xA1, 0x1B, 0xF5, 0x27, 0xC0, 0x56,
        0xDA, 0xFF, 0xD8, 0x2A, 0xB6, 0xCB, 0xC2, 0x5C,
        0xCA, 0x37, 0x14, 0x34, 0x59, 0xE7, 0xBC, 0x63,
        0xBC, 0xDE, 0x52, 0x75, 0x7A, 0xDE, 0xB7, 0xDF,
        0x01, 0xCF, 0x12, 0x17, 0x3F, 0x1F, 0xEF, 0x81,
        0x02, 0xEC, 0x5A, 0xB1, 0x42, 0xC2, 0x13, 0xDD,
        0x9D, 0x30, 0x69, 0x62, 0x78, 0xA8, 0xD8, 0xBC,
        0x32, 0xDD, 0xE9, 0x59, 0x2D, 0x28, 0xC0, 0x78,
        0xC6, 0xD9, 0x2B, 0x94, 0x7D, 0x82, 0x5A, 0xCA,
        0xAB, 0x64, 0x94, 0x84, 0x6A, 0x49, 0xDE, 0x24,
        0xB9, 0x62, 0x3F, 0x48, 0x89, 0xE8, 0xAD, 0xC3,
        0x8E, 0x8C, 0x66, 0x9E, 0xFF, 0xEF, 0x17, 0x60,
        0x40, 0xAD, 0x94, 0x5E, 0x90, 0xA7, 0xD3, 0xEE,
        0xC1, 0x5E, 0xFE, 0xEE, 0x78, 0xAE, 0x71, 0x04,
        0x3C, 0x96, 0x51, 0x11, 0x03, 0xA1, 0x6B, 0xA7,
        0xCA, 0xF0, 0xAC, 0xD0, 0x64, 0x2E, 0xFD, 0xBE,
        0x80, 0x99, 0x34, 0xFA, 0xA1, 0xA5, 0xF1, 0xBD,
        0x11, 0x04, 0x36, 0x49, 0xB2, 0x5C, 0xCD, 0x1F,
        0xEE, 0x2E, 0x38, 0x81, 0x5D, 0x4D, 0x5F, 0x5F,
        0xC6, 0xB4, 0x10, 0x29, 0x69, 0xF2, 0x1C, 0x22,
        0xAE, 0x1B, 0x0E, 0x7D, 0x36, 0x03, 0xA5, 0x56,
        0xA1, 0x32, 0x62, 0xFF, 0x62, 0x8D, 0xE2, 0x22
    ];
    let ssh_kdf_set3_h = [
        0x7B, 0x70, 0x01, 0x18, 0x5E, 0x25, 0x6D, 0x44,
        0x93, 0x44, 0x5F, 0x39, 0xA5, 0x5F, 0xB9, 0x05,
        0xE6, 0x32, 0x1F, 0x4B, 0x5D, 0xD8, 0xBB, 0xF3,
        0x10, 0x0D, 0x51, 0xBA, 0x0B, 0xDA, 0x3D, 0x2D
    ];
    let ssh_kdf_set3_sid = [
        0x7B, 0x70, 0x01, 0x18, 0x5E, 0x25, 0x6D, 0x44,
        0x93, 0x44, 0x5F, 0x39, 0xA5, 0x5F, 0xB9, 0x05,
        0xE6, 0x32, 0x1F, 0x4B, 0x5D, 0xD8, 0xBB, 0xF3,
        0x10, 0x0D, 0x51, 0xBA, 0x0B, 0xDA, 0x3D, 0x2D
    ];
    let ssh_kdf_set3_a = [
        0x81, 0xF0, 0x33, 0x0E, 0xF6, 0xF0, 0x53, 0x61,
        0xB3, 0x82, 0x3B, 0xFD, 0xED, 0x6E, 0x1D, 0xE9
    ];

    let mut out = [0u8; 16];

    ssh_kdf(HMAC::TYPE_SHA256, b'A', &ssh_kdf_set3_k, &ssh_kdf_set3_h,
        &ssh_kdf_set3_sid, &mut out).expect("Error with ssh_kdf()");

    assert_eq!(out, ssh_kdf_set3_a);
}

#[test]
#[cfg(kdf_srtp)]
fn test_srtp_kdf() {
    let key = [
        0xc4u8, 0x80, 0x9f, 0x6d, 0x36, 0x98, 0x88, 0x72,
        0x8e, 0x26, 0xad, 0xb5, 0x32, 0x12, 0x98, 0x90
    ];
    let salt = [
        0x0eu8, 0x23, 0x00, 0x6c, 0x6c, 0x04, 0x4f, 0x56,
        0x62, 0x40, 0x0e, 0x9d, 0x1b, 0xd6
    ];
    let index = [
        0x48u8, 0x71, 0x65, 0x64, 0x9c, 0xca
    ];
    let expected_ke = [
        0xdcu8, 0x38, 0x21, 0x92, 0xab, 0x65, 0x10, 0x8a,
        0x86, 0xb2, 0x59, 0xb6, 0x1b, 0x3a, 0xf4, 0x6f
    ];
    let expected_ka = [
        0xb8u8, 0x39, 0x37, 0xfb, 0x32, 0x17, 0x92, 0xee,
        0x87, 0xb7, 0x88, 0x19, 0x3b, 0xe5, 0xa4, 0xe3,
        0xbd, 0x32, 0x6e, 0xe4
    ];
    let expected_ks = [
        0xf1u8, 0xc0, 0x35, 0xc0, 0x0b, 0x5a, 0x54, 0xa6,
        0x16, 0x92, 0xc0, 0x16, 0x27, 0x6c
    ];
    let mut key_e = [0u8; 16];
    let mut key_a = [0u8; 20];
    let mut key_s = [0u8; 14];
    srtp_kdf(&key, &salt, -1, &index, &mut key_e, &mut key_a, &mut key_s).expect("Error with srtp_kdf()");
    assert_eq!(key_e, expected_ke);
    assert_eq!(key_a, expected_ka);
    assert_eq!(key_s, expected_ks);

    let mut key_e = [0u8; 16];
    srtp_kdf_label(&key, &salt, -1, &index, SRTP_LABEL_ENCRYPTION, &mut key_e).expect("Error with srtp_kdf_label()");
    assert_eq!(key_e, expected_ke);

    let mut key_a = [0u8; 20];
    srtp_kdf_label(&key, &salt, -1, &index, SRTP_LABEL_MSG_AUTH, &mut key_a).expect("Error with srtp_kdf_label()");
    assert_eq!(key_a, expected_ka);

    let mut key_s = [0u8; 14];
    srtp_kdf_label(&key, &salt, -1, &index, SRTP_LABEL_SALT, &mut key_s).expect("Error with srtp_kdf_label()");
    assert_eq!(key_s, expected_ks);
}

#[test]
#[cfg(kdf_srtp)]
fn test_srtcp_kdf() {
    let key = [
        0xc4u8, 0x80, 0x9f, 0x6d, 0x36, 0x98, 0x88, 0x72,
        0x8e, 0x26, 0xad, 0xb5, 0x32, 0x12, 0x98, 0x90
    ];
    let salt = [
        0x0eu8, 0x23, 0x00, 0x6c, 0x6c, 0x04, 0x4f, 0x56,
        0x62, 0x40, 0x0e, 0x9d, 0x1b, 0xd6
    ];
    let index = [
        0x56u8, 0xf3, 0xf1, 0x97
    ];
    let expected_ke = [
        0xabu8, 0x5b, 0xe0, 0xb4, 0x56, 0x23, 0x5d, 0xcf,
        0x77, 0xd5, 0x08, 0x69, 0x29, 0xba, 0xfb, 0x38
    ];
    let expected_ka = [
        0xc5u8, 0x2f, 0xde, 0x0b, 0x80, 0xb0, 0xf0, 0xba,
        0xd8, 0xd1, 0x56, 0x45, 0xcb, 0x86, 0xe7, 0xc7,
        0xc3, 0xd8, 0x77, 0x0e
    ];
    let expected_ks = [
        0xdeu8, 0xb5, 0xf8, 0x5f, 0x81, 0x33, 0x6a, 0x96,
        0x5e, 0xd3, 0x2b, 0xb7, 0xed, 0xe8
    ];
    let mut key_e = [0u8; 16];
    let mut key_a = [0u8; 20];
    let mut key_s = [0u8; 14];
    srtcp_kdf(&key, &salt, -1, &index, &mut key_e, &mut key_a, &mut key_s).expect("Error with srtcp_kdf()");
    assert_eq!(key_e, expected_ke);
    assert_eq!(key_a, expected_ka);
    assert_eq!(key_s, expected_ks);

    let mut key_e = [0u8; 16];
    srtcp_kdf_label(&key, &salt, -1, &index, SRTCP_LABEL_ENCRYPTION, &mut key_e).expect("Error with srtcp_kdf_label()");
    assert_eq!(key_e, expected_ke);

    let mut key_a = [0u8; 20];
    srtcp_kdf_label(&key, &salt, -1, &index, SRTCP_LABEL_MSG_AUTH, &mut key_a).expect("Error with srtcp_kdf_label()");
    assert_eq!(key_a, expected_ka);

    let mut key_s = [0u8; 14];
    srtcp_kdf_label(&key, &salt, -1, &index, SRTCP_LABEL_SALT, &mut key_s).expect("Error with srtcp_kdf_label()");
    assert_eq!(key_s, expected_ks);
}

#[test]
#[cfg(kdf_srtp)]
fn test_srtp_kdr_to_idx() {
    assert_eq!(srtp_kdr_to_index(0), -1);
    assert_eq!(srtp_kdr_to_index(1), 0);
    assert_eq!(srtp_kdr_to_index(2), 1);
    assert_eq!(srtp_kdr_to_index(4), 2);
    assert_eq!(srtp_kdr_to_index(8), 3);
    assert_eq!(srtp_kdr_to_index(16), 4);
    assert_eq!(srtp_kdr_to_index(65536), 16);
    assert_eq!(srtp_kdr_to_index(1048576), 20);
}