nkode-protocol/src/app/client.rs

use std::marker::PhantomData;
use nkode_rs::nkode_core::keypad::Keypad;
use crate::models::app::{OpaqueAPI, AuthAPI, CodeLoginData, CodeLoginSession, Icon, IconID, KeyLoginSession, ICON_ID_SIZE, LoginSession};
use crate::models::email::Email;
use crate::models::opaque::{OpaqueLoginSession, UserSecretKey};
use anyhow::Result;
use nkode_rs::nkode_core::nkode_cipher::NKodeCipher;
use nkode_rs::from_bytes::FromBytes;
use nkode_rs::nkode_core::chacha20prng::Nonce;
use async_trait::async_trait;
use nkode_rs::nkode_core::policy::{NKodePolicy, DEFAULT_POLICY};
use crate::opaque::client::{AuthenticationData, OpaqueAuthLogin, OpaqueAuthRegister, ServerConnectionLogin, ServerConnectionRegister};

pub struct Login;
pub struct Register;


pub struct ClientAppKey<State, S: OpaqueAPI> {
    api: S,
    email: Email,
    user_secret_key: UserSecretKey,
    _state: PhantomData<State>
}

impl <S: OpaqueAPI> ClientAppKey<Register,S> {
    pub async fn register(self) -> Result<ClientAppKey<Login, S>, String> {
        self.api.register_key(&self.email, &self.user_secret_key).await?;
        Ok(ClientAppKey {
            api: self.api,
            email: self.email,
            user_secret_key: self.user_secret_key,
            _state: PhantomData::<Login>,
        })
    }
}

impl <S: OpaqueAPI> ClientAppKey<Login,S> {
    pub async fn login(self) -> Result<ClientAppKeyLoggedIn<S>,String> {
        let key_login = self.api.login_key(&self.email, &self.user_secret_key).await?;
        Ok(ClientAppKeyLoggedIn{
            api: self.api,
            email: self.email,
            user_secret_key: self.user_secret_key,
            key_login
        })
    }
}

pub struct ClientAppKeyLoggedIn<S: OpaqueAPI> {
    api: S,
    email: Email,
    user_secret_key: UserSecretKey,
    key_login: KeyLoginSession,
}

impl <S: OpaqueAPI + AuthAPI> ClientAppKeyLoggedIn<S> {
    pub async fn register_code(self) -> Result<ClientAppCodeRegister<S>, String> {
        let icon_nonce = Nonce::new();
        let icons = self.get_icons(&icon_nonce).await?;
        let policy = self.api.get_policy().await?;
        Ok(ClientAppCodeRegister {
            api: self.api,
            email: self.email,
            user_secret_key: self.user_secret_key,
            key_login: self.key_login,
            icons,
            icon_nonce,
            keypad: Keypad::new(policy),
        })
    }

    async fn get_icons(&self, icon_nonce: &Nonce) -> Result<Vec<Icon>, String> {
        let chacha20_key = self.user_secret_key.chacha20_secret_key();
        let mut chacha_cipher = nkode_rs::nkode_core::chacha20prng::ChaCha20PRNG::new(&chacha20_key, icon_nonce);
        let mut icons = self.api.get_new_icons(&self.key_login).await?;
        for icon in &mut icons {
            let bytes = chacha_cipher.read_u8(ICON_ID_SIZE)?;
            let new_id = IconID::from_bytes(bytes.as_slice()).unwrap();
            icon.update_id(new_id);
        }
        Ok(icons)
    }

    pub async fn login_code(self) -> Result<ClientAppCodeLogin<S>, String> {
        let login_data = self.api.get_login_data(&self.key_login).await?;
        let icons = self.get_icons(&login_data.icon_nonce()).await?;
        let policy = self.api.get_policy().await?;
        let nkode_secret_key = self.user_secret_key.chacha20_secret_key();
        let cipher = NKodeCipher::from_nonce(policy, &nkode_secret_key, login_data.icon_nonce())?;
        Ok(
            ClientAppCodeLogin{
                api: self.api,
                email: self.email,
                mask: login_data.mask,
                icons,
                keypad: login_data.keypad,
                cipher,
            }
        )
    }
}

pub struct ClientAppCodeRegister<S: OpaqueAPI> {
    api: S,
    email: Email,
    user_secret_key: UserSecretKey,
    key_login: KeyLoginSession,
    icon_nonce: Nonce,
    icons: Vec<Icon>,
    keypad: Keypad,
}

impl <S: OpaqueAPI + AuthAPI>ClientAppCodeRegister<S>  {
    pub async fn register(self, selected_icons: Vec<IconID>) -> Result<ClientAppCodeLogin<S>, String> {
        let policy = self.api.get_policy().await?;
        let keypad = Keypad::new(policy.clone());
        let secret_key = self.user_secret_key.chacha20_secret_key();
        let (cipher, cipher_nonce) = NKodeCipher::new(policy, &secret_key);
        let passcode_idx: Vec<usize> = selected_icons
            .iter()
            .filter_map(|x1| self.icons.iter().position(|x| x.id() == x1))
            .collect();
        let ciphered_nkode = cipher.encipher(&passcode_idx).unwrap();
        let data = CodeLoginData{
            mask: ciphered_nkode.mask,
            cipher_nonce,
            icon_nonce: self.icon_nonce,
            keypad,
        };
        self.api.register_code(&self.email, &ciphered_nkode.passcode, &self.key_login, &data).await?;
        Ok(ClientAppCodeLogin{
            api: self.api,
            mask: data.mask,
            email: self.email,
            keypad: data.keypad,
            cipher,
            icons: self.icons
        })
    }
}

pub struct ClientAppCodeLogin<S: OpaqueAPI> {
    api: S,
    email: Email,
    mask: Vec<u64>,
    icons: Vec<Icon>,
    keypad: Keypad,
    cipher: NKodeCipher
}

impl <S: OpaqueAPI>ClientAppCodeLogin<S> {
    pub fn sort_icons(&self) -> Vec<Icon> {
        nkode_rs::tensor::reorder(
            &self.icons,
            self.keypad.indices()
        ).unwrap().to_vec()
    }

    pub async fn login(&self, selected_keys: &Vec<usize>) -> Result<CodeLoginSession, String> {
        let passcode = self.cipher.decipher(selected_keys, self.keypad.indices(), &self.mask).map_err(|e| format!("invalid keys: {e}"))?;
        self.api.login_code(&self.email, &passcode).await
    }
}

pub struct ClientAuth<
    R: ServerConnectionRegister + Clone,
    L: ServerConnectionLogin + Clone,
    K: AuthAPI
> {
    opaque_key_register: OpaqueAuthRegister<R>,
    opaque_key_login: OpaqueAuthLogin<L>,
    opaque_code_register:  OpaqueAuthRegister<R>,
    opaque_code_login: OpaqueAuthLogin<L>,
    nkode_api: K,
}

#[async_trait]
impl<R, L, K> OpaqueAPI for ClientAuth<R, L,K>
where
    R: ServerConnectionRegister + Clone + Sync + Send,
    L: ServerConnectionLogin   + Clone + Sync + Send,
    K: AuthAPI + Sync + Send,
{
    async fn register_key(&self, email: &Email, secret_key: &UserSecretKey) -> Result<(), String> {
        let auth_data = AuthenticationData::from_secret_key(email.as_str(), secret_key.as_slice());
        self.opaque_key_register.register(&auth_data).await.map_err(|e| format!("error: {}", e))
    }

    async fn register_code(&self, email: &Email, passcode: &[u64], key_login_session: &KeyLoginSession, data: &CodeLoginData) -> Result<(), String> {
        let auth_data = AuthenticationData::from_code(email.as_str(), passcode);
        self.opaque_code_register.register(&auth_data).await.map_err(|e| format!("error: {}", e))
    }

    async fn login_key(&self, email: &Email, secret_key: &UserSecretKey) -> Result<KeyLoginSession, String> {
        let auth_data = AuthenticationData::from_secret_key(&email.as_str(), secret_key.as_slice());
        let session_key = self.opaque_key_login.login(&auth_data).await.map_err(|e| format!("error: {}", e))?;
        Ok(KeyLoginSession(
            LoginSession {
                email: email.clone(),
                session_key
            }
        ))
    }

    async fn login_code(&self, email: &Email, passcode: &[u64]) -> Result<CodeLoginSession, String> {
        let auth_data = AuthenticationData::from_code(email.as_str(), passcode);
        let session_key = self.opaque_code_login.login(&auth_data).await.map_err(|e| format!("error: {}", e))?;
        Ok(CodeLoginSession(
            LoginSession {
                email: email.clone(),
                session_key
            }
        ))
    }
}

#[async_trait]
impl<R, L, K> AuthAPI for ClientAuth<R, L, K>
where
    R: ServerConnectionRegister + Clone + Sync + Send,
    L: ServerConnectionLogin   + Clone + Sync + Send,
    K: AuthAPI + Sync + Send,
{
    async fn get_new_icons(
        &self,
        key_login_session: &KeyLoginSession,
    ) -> Result<Vec<Icon>, String> {
        self.nkode_api
            .get_new_icons(key_login_session)
            .await
    }

    async fn get_login_data(
        &self,
        key_login_session: &KeyLoginSession,
    ) -> Result<CodeLoginData, String> {
        self.nkode_api
            .get_login_data(key_login_session)
            .await
    }

    async fn is_code_registered(&self, key_login_session: &KeyLoginSession) -> Result<bool, String> {
        self.nkode_api.is_code_registered(key_login_session).await
    }

    async fn get_policy(&self) -> Result<NKodePolicy, String> {
        Ok(DEFAULT_POLICY)
    }
}