#![doc = include_str!("../README.md")]
//! # Example
//! Retrieving a key from Locky
//! ```rust
//! # use locky::{LockyClient, LockyEnv};
//! # use aes_gcm::{
//! #       aead::{Aead, AeadCore, KeyInit, OsRng},
//! #       Aes256Gcm, Nonce, Key
//! # };
//! # tokio_test::block_on(async {
//! # // make an account for testing
//! # let (account_id, access_token) = locky::get_test_account().await;
//! // Connect to Locky staging environment.
//! let mut client = LockyClient::new(LockyEnv::Staging)
//!     .with_creds(account_id, access_token);
//! # client.create_key("test_db_key").await.unwrap();
//!
//! // Securely get a secret from the cloud service
//! let key = client.get_key("test_db_key").await.unwrap();
//!
//! // Use the secret to encrypt some data
//! let cipher = Aes256Gcm::new((&*key).into());
//!
//! // Never send this key over a network. Even if the communication is encrypted,
//! // unless it specifially uses a post-quantum secure protocol (such as the one
//! // one used by Locky) it will vulnerable to harvest-now decrypt-later
//! // attacks.
//! drop(key);
//! # });
//! ```
//!
//! ## Creating an account
//! ```rust
//! # use locky::{LockyClient, LockyEnv};
//! # tokio_test::block_on(async {
//! let mut client = LockyClient::new(LockyEnv::Staging);
//!
//! // Make an account in our staging environment
//! let account_id = client.create_account("cool-test-account@getloc.ky").await.unwrap();
//!
//! // the access token needs to be stored securely, but it does not need
//! // to be stored in a quantum-secure manner. So however you currently
//! // manage secrets is probably fine!
//! let access_token = client.get_access_token().unwrap();
//! # });
//! ```
//!
//! ## Creating a key
//! ```rust
//! # use locky::{LockyClient, LockyEnv};
//! # tokio_test::block_on(async {
//! # let (account_id, access_token) = locky::get_test_account().await;
//! let mut client = LockyClient::new(LockyEnv::Staging).with_creds(account_id, access_token);
//!
//! // Alternately, you can use our CLI or web interface to create a key
//! client.create_key("test_key").await.unwrap();
//! # });
//! ```
//!
//! # A Note On Staging
//! The staging environment is **deleted every 24 hours**. It is a test environment.
//! Security is not guaranteed and any accounts, keys, or data you create
//! will be lost. Do not store anything in staging besides ephemeral test data!
#![doc(html_logo_url = "https://github.githubassets.com/images/icons/emoji/unicode/1f9e9.png")]
#![doc(html_favicon_url = "https://github.githubassets.com/images/icons/emoji/unicode/1f9e9.png")]
mod autogenerated;
use aes_kw::KekAes256;
use autogenerated::api::locky_client::LockyClient as GrpcLockyClient;
use autogenerated::api::{CreateAccountRequest, CreateKeyRequest, GetKeyRequest};
use ml_kem_rs::ml_kem_768;
use ml_kem_rs::ml_kem_768::{CipherText, DecapsKey};
use tonic::transport::{Channel, ClientTlsConfig};
use zeroize::{Zeroize, Zeroizing};

/// LockyClient is a client for interacting with the Locky service.
///
/// It provides methods for creating an account, managing credentials, creating keys, and retrieving keys.
/// The client can be configured with different environments, such as staging or production.
pub struct LockyClient {
    client: Option<GrpcLockyClient<Channel>>,
    creds: Option<(String, String)>,
    env: LockyEnv,
}

pub enum LockyEnv {
    Staging,
    Production,
}

impl LockyClient {
    pub fn new(env: LockyEnv) -> Self {
        LockyClient {
            client: None,
            creds: None,
            env,
        }
    }

    pub fn with_creds<S>(self, account_id: S, access_token: S) -> Self
    where
        S: Into<String>,
    {
        LockyClient {
            creds: Some((account_id.into(), access_token.into())),
            ..self
        }
    }

    pub async fn create_account<S>(
        &mut self,
        email: S,
    ) -> Result<String, Box<dyn std::error::Error>>
    where
        S: Into<String>,
    {
        match self.creds.as_ref() {
            Some(_) => {
                return Err("already logged in! use a new LockyClient to make a new account".into())
            }
            None => {
                let request = tonic::Request::new(CreateAccountRequest {
                    email: email.into(),
                });
                let response = self
                    .get_client()
                    .await?
                    .create_account(request)
                    .await?
                    .into_inner();
                self.creds = Some((response.account_id.clone(), response.access_token));
                Ok(response.account_id)
            }
        }
    }

    pub fn get_access_token(&self) -> Result<String, Box<dyn std::error::Error>> {
        match self.creds.as_ref() {
            None => Err("must create_account or use with_creds to provide credentials".into()),
            Some((_, access_token)) => Ok(access_token.clone()),
        }
    }

    pub async fn create_key<S>(&mut self, name: S) -> Result<(), Box<dyn std::error::Error>>
    where
        S: Into<String>,
    {
        match self.creds.as_ref() {
            None => {
                return Err("must create_account or use with_creds to provide credentials".into())
            }
            Some((account_id, access_token)) => {
                let request = tonic::Request::new(CreateKeyRequest {
                    account_id: account_id.clone(),
                    access_token: access_token.clone(),
                    name: name.into(),
                });
                self.get_client().await?.create_key(request).await?;
                Ok(())
            }
        }
    }

    pub async fn get_key<S>(
        &mut self,
        name: S,
    ) -> Result<Zeroizing<[u8; 32]>, Box<dyn std::error::Error>>
    where
        S: Into<String>,
    {
        match self.creds.as_ref() {
            None => {
                return Err("must create_account or use with_creds to provide credentials".into())
            }
            Some((account_id, access_token)) => {
                let (ek, dk) = ml_kem_768::key_gen();
                let request = tonic::Request::new(GetKeyRequest {
                    account_id: account_id.clone(),
                    access_token: access_token.clone(),
                    name: name.into(),
                    ephemeral_encaps_key: ek.to_bytes().to_vec(),
                });
                let response = self
                    .get_client()
                    .await?
                    .get_key(request)
                    .await?
                    .into_inner();
                let ct = ml_kem_768::new_ct(
                    response
                        .encaps_ciphertext
                        .try_into()
                        .map_err(|_| "bad encaps_ciphertext")?,
                );
                let key = decrypt_key(
                    &dk,
                    &ct,
                    response
                        .key_ciphertext
                        .try_into()
                        .map_err(|_| "bad key_ciphertext")?,
                )?;
                Ok(key)
            }
        }
    }

    async fn get_client(
        &mut self,
    ) -> Result<&mut GrpcLockyClient<Channel>, Box<dyn std::error::Error>> {
        if self.client.is_none() {
            match self.env {
                LockyEnv::Staging => Ok(self.client.insert({
                    GrpcLockyClient::new(
                        Channel::from_static("https://api.staging.getloc.ky:443")
                            .tls_config(
                                ClientTlsConfig::new().domain_name("api.staging.getloc.ky"),
                            )?
                            .connect()
                            .await?,
                    )
                })),
                LockyEnv::Production => {
                    unimplemented!("Locky production environment not yet supported");
                }
            }
        } else {
            Ok(self.client.as_mut().unwrap())
        }
    }
}

fn decrypt_key(
    dk: &DecapsKey,
    ct: &CipherText,
    mut to_dec: [u8; 40],
) -> Result<Zeroizing<[u8; 32]>, Box<dyn std::error::Error>> {
    let ssk = dk.decaps(&ct);
    let kek = KekAes256::from(ssk.to_bytes());
    let mut res = Zeroizing::new([0u8; 32]);
    kek.unwrap(&to_dec, res.as_mut())
        .map_err(|_| "failed to decrypt key from Locky")?;
    to_dec.zeroize();
    Ok(res)
}

#[doc(hidden)]
pub async fn get_test_account() -> (String, String) {
    let mut client = LockyClient::new(LockyEnv::Staging);
    let account_id = client
        .create_account("doctest-acct@getloc.ky")
        .await
        .unwrap();
    (account_id, client.get_access_token().unwrap())
}