main.rs - 2

use aes_gcm::aead::{Aead, NewAead};
use aes_gcm::{Aes256Gcm, Key, Nonce};
use rand::Rng;
use rand_core::OsRng;
use std::fs;
use std::hash::{DefaultHasher, Hash, Hasher};
use std::io::{self, Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{SystemTime, UNIX_EPOCH};
use x25519_dalek::{EphemeralSecret, PublicKey};

mod config;

use config::CONFIG;

enum Mode {
    Server,
    Client,
}

enum MessageType {
    Text(String),
    File { name: String, content: Vec<u8> },
}

struct ClientInfo {
    stream: TcpStream,
    id: String,
    public_key: PublicKey,
    shared_secret: Option<[u8; 32]>,
}

fn generate_id() -> String {
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_millis();

    let mut hasher = DefaultHasher::new();

    now.hash(&mut hasher);

    format!("{:04x}", hasher.finish() % 65536)
}

fn save_key(key: &[u8], prefix: &str, id: &str) {
    if CONFIG.debug {
        let filename = format!("{}-key#{}.pem", prefix, id);
        fs::write(&filename, key).unwrap();
        println!("Debug: Saved {} to {}", prefix, filename);
    }
}

fn save_file(content: &[u8], filename: &str) {
    if CONFIG.debug {
        fs::write(filename, content).unwrap();
        println!("Debug: Saved unencrypted file to {}", filename);
    }
}

// fn debug_save_ephemeral_secret(secret: &EphemeralSecret, prefix: &str, id: &str) {
//     if CONFIG.debug {
//         let public_key = PublicKey::from(secret);
//         let filename = format!("{}-key#{}.pem", prefix, id);
//         fs::write(&filename, public_key.as_bytes()).unwrap();
//         println!("Debug: Saved {} public key to {}", prefix, filename);
//     }
// }

fn debug_print(message: &str, data: &[u8]) {
    if CONFIG.debug {
        println!("Debug: {} - {:?}", message, data);
        let mut hasher = DefaultHasher::new();
        data.hash(&mut hasher);
        println!("Debug: {} hash - {:x}", message, hasher.finish());
    }
}

fn encrypt_message(message: &str, shared_secret: &[u8; 32]) -> Vec<u8> {
    let key = Key::from_slice(shared_secret);
    let cipher = Aes256Gcm::new(key);
    let pre_nonce = rand::thread_rng().gen::<[u8; 12]>();
    let nonce = Nonce::from_slice(&pre_nonce);
    let ciphertext = cipher.encrypt(nonce, message.as_bytes()).unwrap();

    debug_print("Unencrypted message", message.as_bytes());
    debug_print("Encrypted message", &ciphertext);
    [nonce.to_vec(), ciphertext].concat()
}

fn decrypt_message(encrypted: &[u8], shared_secret: &[u8; 32]) -> String {
    let key = Key::from_slice(shared_secret);
    let cipher = Aes256Gcm::new(key);
    let nonce = Nonce::from_slice(&encrypted[..12]);
    let ciphertext = &encrypted[12..];
    let plaintext = cipher.decrypt(nonce, ciphertext).unwrap();
    debug_print("Encrypted message", encrypted);
    debug_print("Decrypted message", &plaintext);
    String::from_utf8(plaintext).unwrap()
}

fn handle_connection(
    mut client_info: ClientInfo,
    clients: Arc<Mutex<Vec<ClientInfo>>>,
    server_addr: SocketAddr,
) {
    let mut buffer = [0; 1024 * 1024];
    let client_addr = client_info.stream.peer_addr().unwrap();
    let client_id = client_info.id.clone();

    if CONFIG.debug {
        save_key(
            client_info.public_key.as_bytes(),
            "client-public",
            &client_id,
        );
        if let Some(secret) = &client_info.shared_secret {
            save_key(secret, "shared-secret", &client_id);
        }
    }

    loop {
        match client_info.stream.read(&mut buffer) {
            Ok(0) => break,
            Ok(bytes_read) => {
                let encrypted_message = &buffer[..bytes_read];
                if let Some(shared_secret) = client_info.shared_secret {
                    let decrypted_message = decrypt_message(encrypted_message, &shared_secret);

                    let message_type = if decrypted_message.starts_with("FILE:") {
                        let parts: Vec<&str> = decrypted_message
                            .strip_prefix("FILE:")
                            .unwrap()
                            .splitn(2, ':')
                            .collect();
                        if parts.len() == 2 {
                            if CONFIG.debug {
                                save_file(parts[1].as_bytes(), &format!("debug_{}", parts[0]));
                            }
                            MessageType::File {
                                name: parts[0].to_string(),
                                content: parts[1].as_bytes().to_vec(),
                            }
                        } else {
                            MessageType::Text(decrypted_message)
                        }
                    } else {
                        MessageType::Text(decrypted_message)
                    };

                    match &message_type {
                        MessageType::Text(text) => println!("{}", text.trim()),
                        MessageType::File { name, .. } => println!("Received file: {}", name),
                    }

                    let mut clients = clients.lock().unwrap();
                    for client in clients.iter_mut() {
                        if client.id != client_id {
                            if let Some(client_shared_secret) = client.shared_secret {
                                let message_to_send = match &message_type {
                                    MessageType::Text(text) => text.clone(),
                                    MessageType::File { name, content } => {
                                        format!(
                                            "FILE:{}:{}",
                                            name,
                                            String::from_utf8_lossy(content)
                                        )
                                    }
                                };
                                let encrypted_for_client =
                                    encrypt_message(&message_to_send, &client_shared_secret);
                                client.stream.write_all(&encrypted_for_client).unwrap();
                                client.stream.flush().unwrap();
                            }
                        }
                    }
                }
            }
            Err(_) => break,
        }
    }

    let mut clients = clients.lock().unwrap();
    clients.retain(|client| client.id != client_id);
    println!("Client@{}#{} disconnected", client_addr.ip(), client_id);
}

fn run_server(addr: SocketAddr) -> io::Result<()> {
    let listener = TcpListener::bind(addr)?;
    println!("Server listening on {}", addr);
    println!("Your IP address is: {}", addr.ip());

    let clients: Arc<Mutex<Vec<ClientInfo>>> = Arc::new(Mutex::new(Vec::new()));
    let clients_clone = Arc::clone(&clients);

    // Spawn a thread to handle server input
    let server_addr = addr;
    thread::spawn(move || loop {
        let mut input = String::new();
        io::stdin().read_line(&mut input).unwrap();
        let message = format!("Server@{}: {}\n", server_addr.ip(), input.trim());

        let mut clients = clients_clone.lock().unwrap();
        for client in clients.iter_mut() {
            if let Some(shared_secret) = client.shared_secret {
                let encrypted_message = encrypt_message(&message, &shared_secret);
                client.stream.write_all(&encrypted_message).unwrap();
                client.stream.flush().unwrap();
            }
        }
    });

    for stream in listener.incoming() {
        match stream {
            Ok(mut stream) => {
                let client_id = generate_id();
                println!(
                    "New client connected: {:?}#{}",
                    stream.peer_addr()?,
                    client_id
                );

                let mut public_key_bytes = [0u8; 32];
                stream.read_exact(&mut public_key_bytes)?;
                let client_public_key = PublicKey::from(public_key_bytes);

                let server_secret = EphemeralSecret::random_from_rng(OsRng);
                let server_public = PublicKey::from(&server_secret);

                stream.write_all(server_public.as_bytes())?;

                let shared_secret = server_secret.diffie_hellman(&client_public_key);

                if CONFIG.debug {
                    save_key(server_public.as_bytes(), "server-public", &client_id);
                    save_key(shared_secret.as_bytes(), "shared-secret", &client_id);
                }

                let client_info = ClientInfo {
                    stream: stream.try_clone()?,
                    id: client_id.clone(),
                    public_key: client_public_key,
                    shared_secret: Some(*shared_secret.as_bytes()),
                };
                let clients = Arc::clone(&clients);
                clients.lock().unwrap().push(client_info);

                let server_addr = addr;
                let client_info = ClientInfo {
                    stream,
                    id: client_id,
                    public_key: client_public_key,
                    shared_secret: Some(*shared_secret.as_bytes()),
                };
                thread::spawn(move || {
                    handle_connection(client_info, clients, server_addr);
                });
            }
            Err(e) => {
                eprintln!("Error accepting client: {}", e);
            }
        }
    }
    Ok(())
}

fn run_client(server_addr: SocketAddr) -> io::Result<()> {
    let mut stream = TcpStream::connect(server_addr)?;
    println!("Connected to server at {}", server_addr);

    let client_addr = stream.local_addr()?;
    let client_id = generate_id();
    println!("Your address is: {}#{}", client_addr.ip(), client_id);

    let client_secret = EphemeralSecret::random_from_rng(OsRng);
    let client_public = PublicKey::from(&client_secret);

    stream.write_all(client_public.as_bytes())?;

    let mut server_public_key_bytes = [0u8; 32];
    stream.read_exact(&mut server_public_key_bytes)?;
    let server_public_key = PublicKey::from(server_public_key_bytes);

    let shared_secret = client_secret.diffie_hellman(&server_public_key);

    if CONFIG.debug {
        save_key(client_public.as_bytes(), "client-public", &client_id);
        save_key(server_public_key.as_bytes(), "server-public", &client_id);
        save_key(shared_secret.as_bytes(), "shared-secret", &client_id);
    }

    let mut receive_stream = stream.try_clone()?;

    let ss_bytes = *shared_secret.as_bytes();

    thread::spawn(move || {
        let mut buffer = [0; 1024 * 1024];
        loop {
            match receive_stream.read(&mut buffer) {
                Ok(0) => break,
                Ok(bytes_read) => {
                    let encrypted_message = &buffer[..bytes_read];
                    let decrypted_message = decrypt_message(encrypted_message, &ss_bytes);
                    if decrypted_message.starts_with("FILE:") {
                        let parts: Vec<&str> = decrypted_message[5..].splitn(2, ':').collect();
                        if parts.len() == 2 {
                            let filename = parts[0];
                            let content = parts[1].as_bytes();
                            fs::write(filename, content).unwrap();
                            if CONFIG.debug {
                                save_file(content, &format!("debug_received_{}", filename));
                            }
                            println!("Received file: {}", filename);
                        }
                    } else {
                        print!("{}", decrypted_message);
                    }
                    io::stdout().flush().unwrap();
                }
                Err(_) => break,
            }
        }
    });

    loop {
        let mut input = String::new();
        io::stdin().read_line(&mut input)?;
        let input = input.trim();

        if input.starts_with("/file ") {
            let file_path = input[6..].trim();
            if let Ok(content) = fs::read(file_path) {
                let file_name = Path::new(file_path).file_name().unwrap().to_str().unwrap();
                let message = format!("FILE:{}:{}", file_name, String::from_utf8_lossy(&content));
                if CONFIG.debug {
                    save_file(&content, &format!("debug_sent_{}", file_name));
                }
                let encrypted_message = encrypt_message(&message, &ss_bytes);
                stream.write_all(&encrypted_message)?;
                stream.flush()?;
                println!("File sent: {}", file_name);
            } else {
                println!("Failed to read file: {}", file_path);
            }
        } else {
            let message = format!("Client@{}#{}: {}\n", client_addr.ip(), client_id, input);
            let encrypted_message = encrypt_message(&message, &ss_bytes);
            stream.write_all(&encrypted_message)?;
            stream.flush()?;
        }
    }
}

fn get_input(prompt: &str, default: &str) -> String {
    println!("{} (default: {})", prompt, default);
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap();
    let input = input.trim();
    if input.is_empty() {
        default.to_string()
    } else {
        input.to_string()
    }
}

fn main() -> io::Result<()> {
    println!("Choose mode: (1) Server, (2) Client");
    let mode = match get_input("Enter 1 or 2", "1").as_str() {
        "1" => Mode::Server,
        "2" => Mode::Client,
        _ => {
            println!("Invalid choice. Defaulting to Server mode.");
            Mode::Server
        }
    };

    match mode {
        Mode::Server => {
            let ip = get_input("Enter IP to bind to", "0.0.0.0");
            let port = get_input("Enter port to listen on", "8080");
            let addr: SocketAddr = format!("{}:{}", ip, port).parse().unwrap();
            run_server(addr)
        }
        Mode::Client => {
            let ip = get_input("Enter server IP", "127.0.0.1");
            let port = get_input("Enter server port", "8080");
            let addr: SocketAddr = format!("{}:{}", ip, port).parse().unwrap();
            run_client(addr)
        }
    }
}