Skip to main content

mega_security_rs/
server.rs

1use crate::{errors::ServerError, keys::KeysPayload, session::{SessionId, SessionIdEncrypted}, utils::_salt, ClientRegistration};
2
3use std::collections::HashMap;
4
5use sha2::{Digest, Sha256};
6use rand::prelude::*;
7use base64::{engine::general_purpose::URL_SAFE, Engine as _};
8
9#[derive(Debug, Default)]
10pub struct Server {
11    confirm_code: Option<String>,
12    clients_registered: HashMap<String, ClientRegistration>,
13    sessions: Vec<SessionId>,
14}
15
16impl Server {
17    pub fn new(confirm_code: Option<String>) -> Self {
18        Self {
19            confirm_code,
20            ..Default::default()
21        }
22    }
23
24    // TODO: this needs to return a confirmation token
25    pub fn register_client(&mut self, client: &ClientRegistration) -> Result<String, ServerError> {
26        if let Some(_) = self.clients_registered.get(&client.id) {
27            return Err(ServerError::ClientAlreadyRegistred(client.id.to_owned()))
28        }
29
30        self.clients_registered.insert(client.id.to_owned(), client.clone());
31        
32        let mut token = [0u8; 16];
33
34        let mut rng = rand::thread_rng();
35        rng.fill_bytes(&mut token);
36
37        let complete_token = match self.confirm_code.clone() {
38            Some(code) => { 
39                let mut t = Vec::new();
40
41                t.append(&mut code.as_bytes().to_vec());
42                t.append(&mut token.to_vec());
43                t.append(&mut client.id.as_bytes().to_vec());
44
45                t
46            },
47            None => { 
48                let mut t = Vec::new();
49
50                t.append(&mut token.to_vec());
51                t.append(&mut client.id.as_bytes().to_vec());
52
53                t
54            }
55        };
56
57        let b64_token = URL_SAFE.encode(complete_token);
58
59        Ok(b64_token)
60    }
61
62    pub fn get_salt_from_id(&self, id: &str) -> Result<Vec<u8>, ServerError> {
63        let client = if let Some(c) = self.clients_registered.get(id) { c } 
64            else { return Err(ServerError::ClientNotFound(id.to_owned()))};
65        
66        Ok(_salt(id, client.random_number()))
67    }
68
69    pub fn auth_client(&mut self, id: &str, autentication_key: &[u8]) -> Result<(KeysPayload, SessionIdEncrypted), ServerError>{
70        let mut auth_key_hasher = Sha256::new();
71        auth_key_hasher.update(autentication_key);
72        let hashed_auth_key = &auth_key_hasher.finalize()[..16];
73
74        let client = if let Some(c) = self.clients_registered.get(id) { c } 
75            else { return Err(ServerError::ClientNotFound(id.to_owned()))};
76
77        let (enc_session_id, token) = SessionIdEncrypted::new(&client.rsa_public_key);
78
79        self.sessions.push(token);
80
81        if client.hashed_auth_key() == hashed_auth_key { Ok((client.encrypted_keys.clone(), enc_session_id)) } 
82            else { Err(ServerError::AutenticationFailed(id.to_owned())) }
83    } 
84
85    pub fn check_session_id(&self, given_session_id: &SessionId) -> Result<(), ServerError>{
86        if self.sessions.contains(&given_session_id) {
87            Ok(())
88        } else {
89            Err(ServerError::SessionIdNotFound)
90        }
91    }
92}