mega-security-rs 0.3.0

An implementation of Mega's security whitepaper
Documentation
use crate::{errors::ServerError, keys::KeysPayload, session::{SessionId, SessionIdEncrypted}, utils::_salt, ClientRegistration};

use std::collections::HashMap;

use sha2::{Digest, Sha256};
use rand::prelude::*;
use base64::{engine::general_purpose::URL_SAFE, Engine as _};

#[derive(Debug, Default)]
pub struct Server {
    confirm_code: Option<String>,
    clients_registered: HashMap<String, ClientRegistration>,
    sessions: Vec<SessionId>,
}

impl Server {
    pub fn new(confirm_code: Option<String>) -> Self {
        Self {
            confirm_code,
            ..Default::default()
        }
    }

    // TODO: this needs to return a confirmation token
    pub fn register_client(&mut self, client: &ClientRegistration) -> Result<String, ServerError> {
        if let Some(_) = self.clients_registered.get(&client.id) {
            return Err(ServerError::ClientAlreadyRegistred(client.id.to_owned()))
        }

        self.clients_registered.insert(client.id.to_owned(), client.clone());
        
        let mut token = [0u8; 16];

        let mut rng = rand::thread_rng();
        rng.fill_bytes(&mut token);

        let complete_token = match self.confirm_code.clone() {
            Some(code) => { 
                let mut t = Vec::new();

                t.append(&mut code.as_bytes().to_vec());
                t.append(&mut token.to_vec());
                t.append(&mut client.id.as_bytes().to_vec());

                t
            },
            None => { 
                let mut t = Vec::new();

                t.append(&mut token.to_vec());
                t.append(&mut client.id.as_bytes().to_vec());

                t
            }
        };

        let b64_token = URL_SAFE.encode(complete_token);

        Ok(b64_token)
    }

    pub fn get_salt_from_id(&self, id: &str) -> Result<Vec<u8>, ServerError> {
        let client = if let Some(c) = self.clients_registered.get(id) { c } 
            else { return Err(ServerError::ClientNotFound(id.to_owned()))};
        
        Ok(_salt(id, client.random_number()))
    }

    pub fn auth_client(&mut self, id: &str, autentication_key: &[u8]) -> Result<(KeysPayload, SessionIdEncrypted), ServerError>{
        let mut auth_key_hasher = Sha256::new();
        auth_key_hasher.update(autentication_key);
        let hashed_auth_key = &auth_key_hasher.finalize()[..16];

        let client = if let Some(c) = self.clients_registered.get(id) { c } 
            else { return Err(ServerError::ClientNotFound(id.to_owned()))};

        let (enc_session_id, token) = SessionIdEncrypted::new(&client.rsa_public_key);

        self.sessions.push(token);

        if client.hashed_auth_key() == hashed_auth_key { Ok((client.encrypted_keys.clone(), enc_session_id)) } 
            else { Err(ServerError::AutenticationFailed(id.to_owned())) }
    } 

    pub fn check_session_id(&self, given_session_id: &SessionId) -> Result<(), ServerError>{
        if self.sessions.contains(&given_session_id) {
            Ok(())
        } else {
            Err(ServerError::SessionIdNotFound)
        }
    }
}