//
// Copyright 2025 The Project Oak Authors
//
// 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.
//

#[rustfmt::skip]
// Wrapping this section in an unformatted module to control how we explain the imports.
mod intro_import {
    // The core types for working with sessions.
    pub use oak_session::{ClientSession, ServerSession};

    // We will also need to configure the sessions. The types in the next block help with that.
    pub use oak_session::{
        attestation::AttestationType,
        config::SessionConfig,
        handshake::HandshakeType,
        key_extractor::DefaultBindingKeyExtractor,
        session_binding::SignatureBinder,
    };

    // The session will use Public Key evidence, endorsed by a Confidential Space JWT.
    pub use p256::ecdsa::{SigningKey, VerifyingKey};
    // These are used by the server,
    pub use oak_attestation::public_key::{PublicKeyAttester, PublicKeyEndorser};
    // This is the endorsemment format for a public key in a JWT tokens.
    pub use oak_proto_rust::oak::attestation::v1::{ConfidentialSpaceEndorsement, ConfidentialSpaceReferenceValues};

    // The client will use these to verify the server's evidence.
    pub use oak_attestation_verification::EventLogVerifier;
    pub use oak_attestation_gcp::policy_generator::confidential_space_policy_from_reference_values;

    // This should be a unique identifier that client and server agree on.
    pub const ATTESTATION_ID: &str = "c0bbb3a6-2256-4390-a342-507b6aecb7e1";

    // This trait provides the `is_open` method that we use during handshake.
    pub use oak_session::Session;

    // These traits provide an easier-to-use interface over the ClientSession and ServerSession.
    pub use oak_session::channel::{SessionInitializer, SessionChannel};
}

use std::{str::FromStr, sync::Arc};

use intro_import::*;
use oak_time::{clock::FixedClock, Instant};

// In this simple example, we show the basics of passing messages back and forth
// bewteen an Oak ClientSession and an Oak ServerSession. In most real-world
// scenarios, you won't have the server and the client located right next to
// each other like this. For a slightly more realistic example, see
// unattested_pair_split.
fn main() {
    // Prepare the signing key to be used to bind sessions. In a real system, this
    // should be generated by the binary itself, and it should be rotated
    // frequently.
    const SIGNING_KEY: &str = include_str!("data/binding_key.pem");

    // The public part of the signing key is hashed and then signed inside of a JWT.
    // Since this is an example that does not run in GCP, we cannot request a real
    // token. This hardcoded token is representative of a real token, signed with
    // a fake Confidential Space root certificate. You can use http://jwt.io to inspect
    // the contents and verify the signature.
    const ENDORSEMENT: &str = include_str!("data/token.jwt");
    const CSPACE_ROOT: &str = include_str!("data/root_ca_cert.pem");

    // Create a server session configured with self-attestation.
    // The attestation is based on a binding key endorsed inside of a JWT.
    // The binding key is used to sign (bind) the open session channel.
    let binding_key = SigningKey::from_str(SIGNING_KEY).expect("Failed to parse binding key");
    let public_key_attester = PublicKeyAttester::new(VerifyingKey::from(&binding_key));
    let public_key_endorser = PublicKeyEndorser::new(ConfidentialSpaceEndorsement {
        jwt_token: ENDORSEMENT.to_owned(),
        ..Default::default()
    });
    let public_key_binder = SignatureBinder::new(Box::new(binding_key));

    let server_config: SessionConfig =
        SessionConfig::builder(AttestationType::SelfUnidirectional, HandshakeType::NoiseNN)
            .add_self_attester(ATTESTATION_ID.to_owned(), Box::new(public_key_attester))
            .add_self_endorser(ATTESTATION_ID.to_owned(), Box::new(public_key_endorser))
            .add_session_binder(ATTESTATION_ID.to_owned(), Box::new(public_key_binder))
            .build();
    let mut server_session =
        ServerSession::create(server_config).expect("Failed to create server session");

    // Create a client session configured with peer attestation.
    // The setup implicitly adds a binding verifier that relies on the key,
    // which is extracted from the evidence.
    let reference_values = ConfidentialSpaceReferenceValues {
        root_certificate_pem: CSPACE_ROOT.to_owned(),
        r#container_image: None,
    };
    // Normally you would use an endorsed policy where the workload (a container) is
    // signed by the developer and the signature committed to Rekor, using Cosign.
    let policy = confidential_space_policy_from_reference_values(&reference_values)
        .expect("failed to generate policy");

    let attestation_verifier = EventLogVerifier::new(
        vec![Box::new(policy)],
        // Tuesday, 1 July 2025 01:30:00 GMT+01:00
        // This time covers the validity of the root certificate and JWT.
        Arc::new(FixedClock::at_instant(Instant::from_unix_seconds(1751391092))),
    );

    let client_config: SessionConfig =
        SessionConfig::builder(AttestationType::PeerUnidirectional, HandshakeType::NoiseNN)
            .add_peer_verifier_with_key_extractor(
                ATTESTATION_ID.to_string(),
                Box::new(attestation_verifier),
                Box::new(DefaultBindingKeyExtractor {}),
            )
            .build();
    let mut client_session =
        ClientSession::create(client_config).expect("Failed to create client session");

    // To start the communications channel, we need to go through the initialization
    // sequence. This includes attestation exchanges, attestation verifications,
    // session bindings, and crypto handshakes.
    //
    // The following sequence works for initialization sequences in any session
    // configuration.
    //
    // Note that in most real use cases, the sender and receiver side will have a
    // loop like this. See unattested_pair_split for a more realistic example.
    while !client_session.is_open() && !server_session.is_open() {
        if !client_session.is_open() {
            // The client will send the next init request.
            let request = client_session.next_init_message().expect("Failed to get init message");
            // The server will receive the next init request
            // A failure could occur if the client sends the wrong kind of message,
            // for example, sending an encrypted message when the initialization
            // sequence is still in process.
            server_session
                .handle_init_message(request)
                .expect("server failed to handle init request");
        }

        // The server *may* have a response.
        if !server_session.is_open() {
            let init_response =
                server_session.next_init_message().expect("failed to get server init response");

            client_session
                .handle_init_message(init_response)
                .expect("Failed to handle server init response");
        }
    }

    println!("Open!");

    // Send one message client to server.
    let message = "Hello Server";
    println!("Client is writing: {message}");
    let to_server = client_session.encrypt(message).expect("failed to encrypt message");

    // Accept and decrypt the message at the server
    let server_read =
        server_session.decrypt(to_server).expect("failed to decrypt message at server");
    let str_message = String::from_utf8_lossy(server_read.as_slice());
    println!("The server read back: {str_message}");
}
