// Copyright 2023 Google LLC
//
// 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.

#![allow(clippy::unwrap_used)]

use crate::{
    proto_adapter::{IntoAdapter as _, MessageType, ToWrappedMessage as _},
    ukey2_handshake::HandshakeCipher,
    HandshakeImplementation, StateMachine, Ukey2ClientStage1, Ukey2ServerStage1,
};
use crypto_provider::elliptic_curve::{EcdhProvider, EphemeralSecret, PublicKey};
use crypto_provider::p256::P256;
use crypto_provider::x25519::X25519;
use crypto_provider::{CryptoProvider, CryptoRng};
use crypto_provider_default::CryptoProviderImpl;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use sha2::Digest;
use std::collections::hash_set;
use ukey2_proto::protobuf::Message;
use ukey2_proto::ukey2_all_proto::ukey;

type X25519EphemeralSecret =
    <<CryptoProviderImpl as CryptoProvider>::X25519 as EcdhProvider<X25519>>::EphemeralSecret;
type X25519PublicKey =
    <<CryptoProviderImpl as CryptoProvider>::X25519 as EcdhProvider<X25519>>::PublicKey;
type P256EphemeralSecret =
    <<CryptoProviderImpl as CryptoProvider>::P256 as EcdhProvider<P256>>::EphemeralSecret;

#[test]
fn advance_from_init_to_finish_client_test() {
    let mut rng = StdRng::from_entropy();
    let client1 = Ukey2ClientStage1::<CryptoProviderImpl>::from(
        &mut rng,
        "next protocol".to_string(),
        HandshakeImplementation::Spec,
    );

    let secret =
        X25519EphemeralSecret::generate_random(&mut <X25519EphemeralSecret as EphemeralSecret<
            X25519,
        >>::Rng::new());
    let public_key = X25519PublicKey::from_bytes(secret.public_key_bytes().as_ref()).unwrap();
    let random: [u8; 32] = rng.gen();
    let message_data: ukey::Ukey2ServerInit = ukey::Ukey2ServerInit {
        version: Some(1),
        random: Some(random.to_vec()),
        handshake_cipher: Some(ukey::Ukey2HandshakeCipher::CURVE25519_SHA512.into()),
        public_key: Some(public_key.to_bytes().as_ref().to_vec()),
        ..Default::default()
    };

    let _client = client1
        .advance_state(&mut rng, &message_data.to_wrapped_msg().write_to_bytes().unwrap())
        .unwrap();
    // TODO assertions on client state
}

#[test]
fn advance_from_init_to_complete_server_x25519_test() {
    let mut rng = StdRng::from_entropy();
    let mut next_protocols = hash_set::HashSet::new();
    let _ = next_protocols.insert("AES_256_CBC-HMAC_SHA256".to_string());
    let server1 = Ukey2ServerStage1::<CryptoProviderImpl>::from(
        next_protocols,
        HandshakeImplementation::Spec,
    );
    // We construct a ClientInit message for the server to get it into the state to handle
    // ClientFinish messages.
    let secret =
        X25519EphemeralSecret::generate_random(&mut <X25519EphemeralSecret as EphemeralSecret<
            X25519,
        >>::Rng::new());
    let client_finished_msg = {
        let mut msg = ukey::Ukey2ClientFinished::default();
        msg.set_public_key(secret.public_key_bytes().as_ref().to_vec());
        msg.to_wrapped_msg()
    };
    let client_finished_bytes = client_finished_msg.write_to_bytes().unwrap();
    let mut hasher = sha2::Sha512::new();
    hasher.update(&client_finished_bytes);
    let client_finished_hash = hasher.finalize().to_vec();
    let cipher = HandshakeCipher::Curve25519Sha512;
    let client_random = rng.gen::<[u8; 32]>();
    let client_init_framed = {
        let commitment = ukey::ukey2client_init::CipherCommitment {
            handshake_cipher: Some(cipher.as_proto().into()),
            commitment: Some(client_finished_hash),
            ..Default::default()
        };
        let client_init = ukey::Ukey2ClientInit {
            version: Some(1),
            random: Some(client_random.to_vec()),
            cipher_commitments: vec![commitment],
            next_protocol: Some("AES_256_CBC-HMAC_SHA256".to_string()),
            ..Default::default()
        };
        client_init.to_wrapped_msg()
    };
    let server2 =
        server1.advance_state(&mut rng, &client_init_framed.write_to_bytes().unwrap()).unwrap();
    assert!(
        !server2.server_init_msg().windows(client_random.len()).any(|w| w == client_random),
        "Server init msg should not contain the client's random"
    );
    // TODO assertions on server2 state
    // We now move the server to the post-ClientFinished state
    let _server = server2.advance_state(&mut rng, &client_finished_bytes).unwrap();
    // TODO assertions on server state
}

#[test]
fn advance_from_init_to_complete_server_p256_test() {
    let mut rng = StdRng::from_entropy();
    let mut next_protocols = hash_set::HashSet::new();
    let _ = next_protocols.insert("AES_256_CBC-HMAC_SHA256".to_string());
    let server1 = Ukey2ServerStage1::<CryptoProviderImpl>::from(
        next_protocols,
        HandshakeImplementation::Spec,
    );
    // We construct a ClientInit message for the server to get it into the state to handle
    // ClientFinish messages.
    let secret = P256EphemeralSecret::generate_random(
        &mut <P256EphemeralSecret as EphemeralSecret<P256>>::Rng::new(),
    );
    let client_finished_msg = {
        let mut msg = ukey::Ukey2ClientFinished::default();
        msg.set_public_key(secret.public_key_bytes().as_ref().to_vec());
        msg.to_wrapped_msg()
    };
    let client_finished_bytes = client_finished_msg.write_to_bytes().unwrap();
    let mut hasher = sha2::Sha512::new();
    hasher.update(&client_finished_bytes);
    let client_finished_hash = hasher.finalize().to_vec();
    let cipher = HandshakeCipher::P256Sha512;
    let client_init_framed = {
        let commitment = ukey::ukey2client_init::CipherCommitment {
            handshake_cipher: Some(cipher.as_proto().into()),
            commitment: Some(client_finished_hash),
            ..Default::default()
        };
        ukey::Ukey2ClientInit {
            version: Some(1),
            random: Some(rng.gen::<[u8; 32]>().to_vec()),
            cipher_commitments: vec![commitment],
            next_protocol: Some("AES_256_CBC-HMAC_SHA256".to_string()),
            ..Default::default()
        }
        .to_wrapped_msg()
    };
    let server2 =
        server1.advance_state(&mut rng, &client_init_framed.write_to_bytes().unwrap()).unwrap();
    // TODO assertions on server2 state
    let _server =
        server2.advance_state(&mut rng, &client_finished_msg.write_to_bytes().unwrap()).unwrap();
    // TODO assertions on server state
}

#[test]
fn message_type_discriminant() {
    assert_eq!(1, ukey::ukey2message::Type::ALERT as i32);
    assert_eq!(2, ukey::ukey2message::Type::CLIENT_INIT as i32);
    assert_eq!(3, ukey::ukey2message::Type::SERVER_INIT as i32);
    assert_eq!(4, ukey::ukey2message::Type::CLIENT_FINISH as i32);
}

#[test]
fn cipher_type_discriminant() {
    assert_eq!(100, ukey::Ukey2HandshakeCipher::P256_SHA512 as i32);
    assert_eq!(200, ukey::Ukey2HandshakeCipher::CURVE25519_SHA512 as i32);
}

#[test]
fn convert_to_message_type() {
    assert_eq!(
        MessageType::ClientInit,
        ukey::ukey2message::Type::CLIENT_INIT.into_adapter().unwrap()
    );
    assert_eq!(
        MessageType::ServerInit,
        ukey::ukey2message::Type::SERVER_INIT.into_adapter().unwrap()
    );
    assert_eq!(
        MessageType::ClientFinish,
        ukey::ukey2message::Type::CLIENT_FINISH.into_adapter().unwrap()
    );
}

#[test]
fn convert_to_cipher_type() {
    assert_eq!(HandshakeCipher::P256Sha512, 100.into_adapter().unwrap());
    assert_eq!(HandshakeCipher::Curve25519Sha512, 200.into_adapter().unwrap());
}
