AIRA/src/main.rs

mod identity;
mod crypto;
mod session_manager;
mod utils;
mod ui_interface;
mod constants;
mod discovery;

use std::{collections::HashMap, env, fmt::Display, fs, io, net::{TcpStream, TcpListener}, path::Path, str::FromStr, sync::{RwLock, Mutex, Arc}, thread, thread::sleep, time::Duration};
use actix_web::{http::{header, CookieBuilder}, App, HttpRequest, HttpResponse, HttpMessage, HttpServer, web, web::Data};
use rand_8::{RngCore, rngs::OsRng};
use session_manager::SessionError;
use tungstenite::{
    WebSocket,
    accept,
};
use serde::{Deserialize, Serialize};
use utils::escape_double_quote;
use uuid::Uuid;
use crate::identity::Identity;
use crate::session_manager::{SessionManager, protocol};
use crate::ui_interface::UiConnection;

fn start_websocket_server(global_vars: &Arc<RwLock<GlobalVars>>) -> u16 {
    let websocket_bind_addr = env::var("AIRA_WEBSOCKET_ADDR").unwrap_or("127.0.0.1".to_owned());
    let websocket_port = env::var("AIRA_WEBSOCKET_PORT").unwrap_or("0".to_owned());
    let server = TcpListener::bind(websocket_bind_addr+":"+&websocket_port).unwrap();
    let websocket_port = server.local_addr().unwrap().port();
    let global_vars_clone = Arc::clone(global_vars);
    thread::spawn(move || {
        for stream in server.incoming() {
            let stream = stream.unwrap();
            stream.set_read_timeout(Some(Duration::from_millis(100))).unwrap();
            match accept(stream) {
                Ok(websocket) => {
                    websocket_worker(websocket, &global_vars_clone);
                }
                Err(e) => print_error(e)
            }
        }
    });
    websocket_port
}

fn on_connected(index: usize, session_manager: &Arc<SessionManager>, ui_connection: &Arc<Mutex<UiConnection>>){
    let mut ui_connection = ui_connection.lock().unwrap();
    ui_connection.on_connected(index);
    session_manager.handle_new_session(&index, ui_connection);
}

fn handle_connect(ip: &str, session_manager: &Arc<SessionManager>, ui_connection: &Arc<Mutex<UiConnection>>){
    match session_manager.connect_to(ip){
        Ok(index) => {
            on_connected(index, session_manager, ui_connection);
        }
        Err(e) => {
            if e != SessionError::AlreadyConnected && e != SessionError::IsUs {
                print_error(e)
            }
        }
    }
}

fn websocket_worker(websocket: WebSocket<TcpStream>, global_vars: &Arc<RwLock<GlobalVars>>){
    let ui_connection = Arc::new(Mutex::new(UiConnection::new(websocket)));
    let global_vars_read = global_vars.read().unwrap();
    global_vars_read.session_manager.set_ui_connection(&ui_connection);
    if global_vars_read.is_backend_running { //ui reconnection
        global_vars_read.session_manager.list_sessions().into_iter().for_each(|index|{
            on_connected(index, &global_vars_read.session_manager, &ui_connection);
        });
    } else {
        match SessionManager::start_listener(&global_vars_read.session_manager) {
            Ok(_) => {}
            Err(e) => println!("{}. You won't be able to receive incomming connections from other peers.", e)
        }
        SessionManager::start_receiver_loop(&global_vars_read.session_manager);
        discovery::advertise_me();
        drop(global_vars_read);
        global_vars.write().unwrap().is_backend_running = true;
    }
    let global_vars_read = global_vars.read().unwrap();
    global_vars_read.session_manager.list_contacts().into_iter().for_each(|contact|{
        let mut ui_connection_locked = ui_connection.lock().unwrap();
        ui_connection_locked.set_as_contact(contact.0, &contact.1, contact.2);
        match global_vars_read.session_manager.load_msgs(&contact.0) {
            Some(msgs) => {
                ui_connection_locked.load_msgs(&contact.0, msgs);
            }
            None => {}
        }
    });
    global_vars_read.session_manager.get_saved_msgs().into_iter().for_each(|entry| {
        ui_connection.lock().unwrap().on_received(&entry.0, &entry.1);
    });
    let session_manager_clone = Arc::clone(&global_vars_read.session_manager);
    let ui_connection_clone = Arc::clone(&ui_connection);
    thread::spawn(move || {
        discovery::discover_peers(|discovery_manager, ip| {
            println!("New peer discovered: {}", ip);
            if ui_connection_clone.lock().unwrap().is_valid {
                handle_connect(&ip, &session_manager_clone, &ui_connection_clone);
            } else {
                discovery_manager.stop_service_discovery();
            }
        });
    });
    loop {
        let mut ui_connection_locked = ui_connection.lock().unwrap();
        let msg_result = ui_connection_locked.read_message();
        drop(ui_connection_locked); //release mutex
        match msg_result {
            Ok(msg) => {
                let args: Vec<&str> = msg.split(" ").collect();
                match args[0] {
                    "connect" => {
                        handle_connect(args[1], &global_vars_read.session_manager, &ui_connection);
                    }
                    "send" => {
                        let buffer = protocol::new_message(args[2..].join(" "));
                        let index: usize = args[1].parse().unwrap();
                        match global_vars_read.session_manager.send_to(&index, &buffer){
                            Ok(_) => {
                                if global_vars_read.session_manager.is_contact(&index) {
                                    match global_vars_read.session_manager.store_msg(&index, true, &buffer) {
                                        Ok(_) => {},
                                        Err(e) => print_error(e)
                                    }
                                }
                            }
                            Err(e) => print_error(e)
                        }
                    }
                    "contact" => {
                        let index: usize = args[1].parse().unwrap();
                        match global_vars_read.session_manager.add_contact(index, args[2..].join(" ")) {
                            Ok(_) => {},
                            Err(e) => print_error(e)
                        }
                    }
                    "uncontact" => {
                        let index: usize = args[1].parse().unwrap();
                        match global_vars_read.session_manager.remove_contact(&index) {
                            Ok(_) => {},
                            Err(e) => print_error(e)
                        }
                    }
                    "fingerprints" => {
                        let index: usize = args[1].parse().unwrap();
                        let (local, peer) = global_vars_read.session_manager.get_public_keys(&index);
                        let local = crypto::generate_fingerprint(&local);
                        let peer = crypto::generate_fingerprint(&peer);
                        ui_connection.lock().unwrap().fingerprints(&local, &peer);
                    }
                    "verify" => {
                        let index: usize = args[1].parse().unwrap();
                        match global_vars_read.session_manager.set_verified(&index) {
                            Ok(_) => {},
                            Err(e) => print_error(e)
                        }
                    }
                    "file" => {
                        let index: usize = args[1].parse().unwrap();
                        let file_name = args[2..].join(" ");
                        let mut ui_connection_locked = ui_connection.lock().unwrap();
                        let buffer = ui_connection_locked.read_binary();
                        drop(ui_connection_locked);
                        match buffer {
                            Ok(buffer) => {
                                match global_vars_read.session_manager.send_to(&index,  &protocol::file(&file_name, &buffer)) {
                                    Ok(_) => {
                                        if global_vars_read.session_manager.is_contact(&index) {
                                            match global_vars_read.session_manager.store_file(&index, &buffer) {
                                                Ok(file_uuid) => {
                                                    ui_connection.lock().unwrap().file_sent(index, &file_name, Some(file_uuid.to_string()));
                                                    match global_vars_read.session_manager.store_msg(&index, true, &[&[protocol::Headers::FILE][..], file_uuid.as_bytes(), &file_name.into_bytes()].concat()) {
                                                        Ok(_) => {}
                                                        Err(e) => print_error(e)
                                                    }
                                                }
                                                Err(e) => print_error(e)
                                            }
                                        } else {
                                            ui_connection.lock().unwrap().file_sent(index, &file_name, None);
                                        }
                                    }
                                    Err(e) => print_error(e)
                                }
                            }
                            Err(e) => print_error(e)
                        }
                    }
                    _ => println!("Unknown: {}", msg)
                }
            }
            Err(e) => {
                match e {
                    tungstenite::Error::Io(ref e) => {
                        if e.kind() != io::ErrorKind::WouldBlock {
                            print_error(e)
                        }
                    }
                    tungstenite::Error::ConnectionClosed => {
                        println!("Websocket connection closed");
                        ui_connection.lock().unwrap().is_valid = false;
                        return
                    }
                    _ => print_error(e)
                }
            }
        }
        sleep(Duration::from_millis(constants::MUTEX_RELEASE_DELAY_MS)); //pause let other threads use ui_connection
    }
}

#[derive(Deserialize, Serialize, Debug)]
struct FileInfo {
    uuid: String,
    file_name: String,
}

fn handle_load_file(req: HttpRequest, file_info: web::Query<FileInfo>) -> HttpResponse {
    println!("Load file request from: {}", req.peer_addr().unwrap());
    match req.cookie(constants::HTTP_COOKIE_NAME) {
        Some(cookie) => {
            let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
            let global_vars = global_vars.read().unwrap();
            match global_vars.http_session_manager.get_name_from_cookie(cookie.value()) {
                Some(_) => {
                    match Uuid::from_str(&file_info.uuid) {
                        Ok(uuid) => {
                            match global_vars.session_manager.load_file(uuid) {
                                Some(buffer) => {
                                    return HttpResponse::Ok().header("Content-Disposition", format!("attachment; filename=\"{}\"", escape_double_quote(html_escape::decode_html_entities(&file_info.file_name).to_string()))).content_type("application/octet-stream").body(buffer);
                                }
                                None => {}
                            }
                        }
                        Err(e) => print_error(e)
                    }
                }
                None => {}
            }
        }
        None => {}
    }
    HttpResponse::NotFound().finish()
}

fn handle_logout(req: HttpRequest) -> HttpResponse {
    println!("Logout request from: {}", req.peer_addr().unwrap());
    match req.cookie(constants::HTTP_COOKIE_NAME) {
        Some(cookie) => {
            let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
            let mut global_vars_write = global_vars.write().unwrap();
            global_vars_write.http_session_manager.remove(cookie.value());
            global_vars_write.session_manager.stop();
            global_vars_write.is_backend_running = false;
            HttpResponse::Found().header(header::LOCATION, "/").finish()
        }
        None => HttpResponse::Unauthorized().finish()
    }
}

fn login(identity: Identity, global_vars: &Arc<RwLock<GlobalVars>>) -> HttpResponse {
    let mut global_vars = global_vars.write().unwrap();
    let cookie_value = global_vars.http_session_manager.register(identity.name.to_string());
    if match global_vars.session_manager.get_identity_uuid() {
        Some(uuid) => uuid != identity.uuid,
        None => true
    } {
        global_vars.session_manager.set_identity(Some(identity));
    }
    let cookie = CookieBuilder::new(constants::HTTP_COOKIE_NAME, cookie_value)
        .http_only(true)
        .max_age(time::Duration::hours(4)
    ).finish();
    HttpResponse::Found().header(header::LOCATION, "/").set_header(header::SET_COOKIE, cookie.to_string()).finish()
}

fn handle_login(req: HttpRequest, params: web::Form<LoginParams>) -> HttpResponse {
    println!("Login request from: {}", req.peer_addr().unwrap());
    match Identity::get_identity(&params.uuid, &params.password) {
        Ok(identity) => {
            let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
            match identity.clear_temporary_files() {
                Ok(n) => println!("Deleted {} files", n),
                Err(e) => print_error(e)
            }
            login(identity, global_vars.get_ref())
        }
        Err(e) => {
            let body = get_login_body(Some(&e.to_string()));
            HttpResponse::NotFound().body(body)
        }
    }
}

fn get_login_body(error_msg: Option<&str>) -> String {
    let mut formated_identities = String::new();
    match Identity::get_identities(){
        Ok(identities) => {
            for i in 0..identities.len() {
                let uuid = identities[i].0.to_string();
                formated_identities += &format!("[\"{}\", \"{}\"]", uuid, identities[i].1);
                if i != identities.len()-1 {
                    formated_identities += ", ";
                }
            }
        }
        Err(e) => {
            print_error(e)
        }
    };
    include_str!("frontend/login.html")
        .replace("IDENTITIES", &formated_identities).replace("ERROR_MSG", &match error_msg {
            Some(error_msg) => format!("Error: {}.", error_msg),
            None => String::new()
        })
}

fn handle_create(req: HttpRequest, params: web::Form<CreateParams>) -> HttpResponse {
    println!("Create request from: {}", req.peer_addr().unwrap());
    if params.password == params.password_confirm {
        match Identity::create_new_identidy(&params.name, &params.password) {
            Ok(identity) => {
                let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
                login(identity, global_vars.get_ref())
            }
            Err(e) => {
                print_error(&e);
                let body = get_login_body(Some(&e.to_string()));
                HttpResponse::Ok().body(body)
            }
        }
    } else {
        let body = get_login_body(Some("Passwords don't match"));
        HttpResponse::Ok().body(body)
    }
}

fn index(req: HttpRequest) -> HttpResponse {
    println!("GET request from: {}", req.peer_addr().unwrap());
    let body = match req.cookie(constants::HTTP_COOKIE_NAME) {
        Some(cookie) => {
            let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
            let global_vars = global_vars.write().unwrap();
            match global_vars.http_session_manager.get_name_from_cookie(cookie.value()) {
                Some(name) => {
                    include_str!("frontend/index.html")
                        .replace("WEBSOCKET_PORT", &global_vars.websocket_port.to_string())
                        .replace("IDENTITY_NAME", &name)
                }
                None => get_login_body(None)
            }
        }
        None => get_login_body(None)
    };
    HttpResponse::Ok().body(body)
}

const JS_CONTENT_TYPE: &str = "text/javascript";

fn handle_static(req: HttpRequest) -> HttpResponse {
    println!("Static: {}", req.path());
    let splits: Vec<&str> = req.path()[1..].split("/").collect();
    if splits[0] == "static" {
        let mut response_builder = HttpResponse::Ok();
        match splits[1] {
            "imgs" => {
                if splits[2] == "icons" && splits.len() <= 5 {
                    let color = if splits.len() == 5 {
                        splits[4]
                    } else {
                        "none"
                    };
                    match match splits[3] {
                        "verified" => Some(include_str!("frontend/imgs/icons/verified.svg")),
                        "add_contact" => Some(include_str!("frontend/imgs/icons/add_contact.svg")),
                        "remove_contact" => Some(include_str!("frontend/imgs/icons/remove_contact.svg")),
                        "logout" => Some(include_str!("frontend/imgs/icons/logout.svg")),
                        "warning" => Some(include_str!("frontend/imgs/icons/warning.svg")),
                        "attach" => Some(include_str!("frontend/imgs/icons/attach.svg")),
                        "download" => Some(include_str!("frontend/imgs/icons/download.svg")),
                        _ => None
                    } {
                        Some(body) => {
                            response_builder.content_type("image/svg+xml");
                            return response_builder.body(body.replace("FILL_COLOR", color))
                        }
                        None => {}
                    }
                } else if splits.len() == 3 {
                    response_builder.content_type("image/jpeg");
                    match splits[2] {
                        "wallpaper" => return response_builder.body(&include_bytes!("frontend/imgs/wallpaper.jpeg")[..]),
                        _ => {}
                    }
                }
            }
            "commons" => {
                if splits.len() == 3 {
                    match splits[2] {
                        "script.js" => return response_builder.content_type(JS_CONTENT_TYPE).body(include_str!("frontend/commons/script.js")),
                        "style.css" => return response_builder.body(include_str!("frontend/commons/style.css")),
                        _ => {}
                    }
                }
            }
            "libs" => {
                if splits.len() == 3 {
                    match match splits[2] {
                        "linkify.min.js" => Some(include_str!("frontend/libs/linkify.min.js")),
                        "linkify-html.min.js" => Some(include_str!("frontend/libs/linkify-html.min.js")),
                        "linkify-element.min.js" => Some(include_str!("frontend/libs/linkify-element.min.js")),
                        _ => None
                    } {
                        Some(body) => return response_builder.content_type(JS_CONTENT_TYPE).body(body),
                        None => {}
                    }
                }
            }
            _ => {}
        }
    }
    HttpResponse::NotFound().finish()
}

#[actix_web::main]
async fn start_http_server(global_vars: Arc<RwLock<GlobalVars>>) -> io::Result<()> {
    let http_bind_addr = env::var("AIRA_HTTP_ADDR").unwrap_or("127.0.0.1".to_owned());
    let http_port = env::var("AIRA_HTTP_PORT").unwrap_or("0".to_owned());
    let server = HttpServer::new(move || {
        let global_vars_clone = global_vars.clone();
        App::new()
            .data(global_vars_clone)
            .service(web::resource("/")
                .route(web::get().to(index))
                .route(web::post().to(handle_create))
            )
            .route("/login", web::post().to(handle_login))
            .route("/load_file", web::get().to(handle_load_file))
            .route("/static/.*", web::get().to(handle_static))
            .route("/logout", web::get().to(handle_logout))
        }
    ).bind(http_bind_addr+":"+&http_port)?;
    let url = "http://127.0.0.1:".to_owned()+&server.addrs().get(0).unwrap().port().to_string();
    println!("AIRA started on: {}", url);
    if webbrowser::open(&url).is_err() {
        println!("Failed to open browser. Please open the link manually.");
    }
    server.run().await
}

#[derive(Serialize, Deserialize)]
struct LoginParams {
    uuid: String,
    password: String
}

#[derive(Serialize, Deserialize)]
struct CreateParams {
    name: String,
    password: String,
    password_confirm: String
}

struct HttpSessionsManager {
    http_sessions: HashMap<String, String>
}

impl HttpSessionsManager {
    pub fn new() -> HttpSessionsManager {
        HttpSessionsManager {
            http_sessions: HashMap::new()
        }
    }
    pub fn register(&mut self, identity_name: String) -> String {
        let mut raw_cookie = [0; 32];
        OsRng.fill_bytes(&mut raw_cookie);
        let cookie = base64::encode(raw_cookie);
        self.http_sessions.insert(cookie.clone(), identity_name);
        cookie
    }
    pub fn remove(&mut self, cookie: &str) {
        self.http_sessions.remove(cookie);
    }
    pub fn get_name_from_cookie(&self, cookie_value: &str) -> Option<String> {
        Some(self.http_sessions.get(cookie_value)?.clone())
    }
}

/*impl Debug for HttpSessionsManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("{\n")?;
        for (k, v) in self.http_sessions.read().unwrap().iter() {
            f.write_str(&format!("\t{} : {}", k, v))?;
        }
        f.write_str("\n}")
        
    }
}*/

struct GlobalVars {
    session_manager: Arc<SessionManager>,
    is_backend_running: bool,
    websocket_port: u16,
    http_session_manager: HttpSessionsManager,
}

fn main() {
    match fs::create_dir(Path::new(&dirs::config_dir().unwrap()).join(constants::APPLICATION_FOLDER)){
        Ok(_) => {}
        Err(e) => {
            if e.kind() != io::ErrorKind::AlreadyExists {
                print_error(e);
            }
        }
    }
    let session_manager = Arc::new(SessionManager::new());
    let global_vars = Arc::new(RwLock::new(GlobalVars {
        session_manager: session_manager,
        is_backend_running: false,
        websocket_port: 0,
        http_session_manager: HttpSessionsManager::new()
    }));
    let websocket_port = start_websocket_server(&global_vars);
    global_vars.write().unwrap().websocket_port = websocket_port;
    start_http_server(global_vars).unwrap();
}

fn print_error<T: Display>(error: T){
    println!("{}", error);
}