nkode-protocol/src/client/states.rs

use std::marker::PhantomData;
use nkode_rs::nkode_core::chacha20prng::Nonce;
use nkode_rs::nkode_core::keypad::Keypad;
use nkode_rs::nkode_core::nkode_cipher::NKodeCipher;
use nkode_rs::from_bytes::FromBytes;
use crate::shared::models::app::{AuthAPI, CodeLoggedInSession, CodeLoginData, Icon, IconID, KeyLoggedInSession, ICON_ID_SIZE};
use crate::shared::email::Email;
use crate::shared::opaque::UserSecretKey;

struct KeyLogin;

struct KeyRegister;

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

impl <S: AuthAPI> UserStateKey<KeyRegister,S> {
    pub async fn register(self) -> Result<UserStateKey<KeyLogin, S>, String> {
        self.api.register_key(&self.email, &self.user_secret_key).await?;
        Ok(UserStateKey {
            api: self.api,
            email: self.email,
            user_secret_key: self.user_secret_key,
            _state: PhantomData::<KeyLogin>,
        })
    }
}

impl <S: AuthAPI> UserStateKey<KeyLogin,S> {
    pub async fn login(self) -> Result<UserStateKeyLoggedIn<S>,String> {
        let key_login = self.api.login_key(&self.email, &self.user_secret_key).await?;
        Ok(UserStateKeyLoggedIn {
            api: self.api,
            email: self.email,
            user_secret_key: self.user_secret_key,
            key_login
        })
    }
}

pub struct UserStateKeyLoggedIn<S: AuthAPI> {
    api: S,
    email: Email,
    user_secret_key: UserSecretKey,
    key_login: KeyLoggedInSession,
}

impl <S: AuthAPI> UserStateKeyLoggedIn<S> {
    pub async fn register_code(self) -> Result<UserStateCodeRegister<S>, String> {
        let icon_nonce = Nonce::new();
        let icons = self.get_icons(&icon_nonce).await?;
        let policy = self.api.get_policy().await?;
        Ok(UserStateCodeRegister {
            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().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<UserStateCodeLogin<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(
            UserStateCodeLogin {
                api: self.api,
                email: self.email,
                mask: login_data.mask,
                icons,
                keypad: login_data.keypad,
                cipher,
            }
        )
    }
}

pub struct UserStateCodeRegister<S: AuthAPI> {
    api: S,
    email: Email,
    user_secret_key: UserSecretKey,
    key_login: KeyLoggedInSession,
    icon_nonce: Nonce,
    icons: Vec<Icon>,
    keypad: Keypad,
}

impl <S: AuthAPI> UserStateCodeRegister<S>  {
    pub async fn register(self, selected_icons: Vec<IconID>) -> Result<UserStateCodeLogin<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.clone()).await?;
        Ok(UserStateCodeLogin {
            api: self.api,
            mask: data.mask,
            email: self.email,
            keypad: data.keypad,
            cipher,
            icons: self.icons
        })
    }
}

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

impl <S: AuthAPI> UserStateCodeLogin<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<CodeLoggedInSession, 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
    }
}