add authentication barebone

3 files changed, 214 insertions, 3 deletions

diff --git a/src/bin/client/connection.rs b/src/bin/client/connection.rs
index d03dc1a..497ed36 100644
--- a/src/bin/client/connection.rs
+++ b/src/bin/client/connection.rs
@@ -3,8 +3,18 @@ use std::net::TcpStream;
 use std::sync::mpsc;
 use std::thread;
 use std::time::{Duration, Instant};
+use rusqlite::Connection;
+use bcrypt::hash_with_salt;
+use crate::database;
 
-pub fn connect(addr: &str) -> Result<TcpStream, Box<dyn std::error::Error>> {
+const SALT: [u8; 16] = *b"fixedsalt1234567";
+
+fn hash_password(password: &str) -> String {
+    let result = hash_with_salt(password, 10, SALT).unwrap();
+    result.to_string()
+}
+
+pub fn connect(addr: &str, conn: &Connection) -> Result<TcpStream, Box<dyn std::error::Error>> {
     let (tx, rx) = mpsc::channel();
     let addr_str = addr.to_string();
     thread::spawn(move || {
@@ -23,6 +33,21 @@ pub fn connect(addr: &str) -> Result<TcpStream, Box<dyn std::error::Error>> {
             ping_reader.read_line(&mut response).unwrap();
             let elapsed = start.elapsed();
             println!("\n\rConnected to {} with ping delay: {:?}", addr, elapsed);
+            
+            // Check for credentials and auto-auth
+            if let Some((username, password)) = database::get_credentials(conn, addr) {
+                let hashed = hash_password(&password);
+                let auth_msg = format!("AUTH {} {}", username, hashed);
+                writeln!(ping_writer, "{}", auth_msg).unwrap();
+                let mut auth_response = String::new();
+                ping_reader.read_line(&mut auth_response).unwrap();
+                if auth_response.trim().starts_with("OK") {
+                    println!("Authenticated as {}", username);
+                } else {
+                    println!("Authentication failed: {}", auth_response.trim());
+                    // Do not fail, allow manual login
+                }
+            }
             Ok(s)
         }
         Ok(Err(e)) => Err(Box::new(e)),
diff --git a/src/bin/client/database.rs b/src/bin/client/database.rs
index bc8f6f8..03a6c45 100644
--- a/src/bin/client/database.rs
+++ b/src/bin/client/database.rs
@@ -1,4 +1,10 @@
 use rusqlite::Connection;
+use aes_gcm::{Aes256Gcm, Key, Nonce};
+use aes_gcm::aead::{Aead, KeyInit};
+use pbkdf2::pbkdf2;
+use hmac::Hmac;
+use sha2::Sha256;
+use rand::Rng;
 
 pub fn init_db() -> Connection {
     let conn = Connection::open("client.db").unwrap();
@@ -9,6 +15,14 @@ pub fn init_db() -> Connection {
         )",
         [],
     ).unwrap();
+    conn.execute(
+        "CREATE TABLE IF NOT EXISTS credentials (
+            host TEXT PRIMARY KEY,
+            username TEXT NOT NULL,
+            encrypted_password BLOB NOT NULL
+        )",
+        [],
+    ).unwrap();
     conn
 }
 
@@ -28,3 +42,53 @@ pub fn save_host(conn: &Connection, host: &str) {
         [host],
     ).unwrap();
 }
+
+pub fn save_credentials(conn: &Connection, host: &str, username: &str, password: &str) -> Result<(), aes_gcm::Error> {
+    let encrypted = encrypt_password(password)?;
+    conn.execute(
+        "INSERT OR REPLACE INTO credentials (host, username, encrypted_password) VALUES (?1, ?2, ?3)",
+        (host, username, &encrypted),
+    ).unwrap();
+    Ok(())
+}
+
+pub fn get_credentials(conn: &Connection, host: &str) -> Option<(String, String)> {
+    let mut stmt = conn.prepare("SELECT username, encrypted_password FROM credentials WHERE host = ?1").unwrap();
+    let mut rows = stmt.query([host]).unwrap();
+    if let Ok(Some(row)) = rows.next() {
+        let username: String = row.get(0).unwrap();
+        let encrypted: Vec<u8> = row.get(1).unwrap();
+        decrypt_password(&encrypted).map(|password| (username, password))
+    } else {
+        None
+    }
+}
+
+fn derive_key() -> Key<Aes256Gcm> {
+    let mut key = [0u8; 32];
+    let salt = b"client_salt"; // Fixed salt for simplicity
+    pbkdf2::<Hmac<Sha256>>(b"client_key", salt, 10000, &mut key).expect("error deriving key");
+    key.into()
+}
+
+fn encrypt_password(password: &str) -> Result<Vec<u8>, aes_gcm::Error> {
+    let key = derive_key();
+    let cipher = Aes256Gcm::new(&key);
+    let mut nonce_bytes = [0u8; 12];
+    rand::thread_rng().fill(&mut nonce_bytes);
+    let nonce = Nonce::from_slice(&nonce_bytes);
+    let ciphertext = cipher.encrypt(nonce, password.as_bytes())?;
+    let mut result = nonce_bytes.to_vec();
+    result.extend(ciphertext);
+    Ok(result)
+}
+
+fn decrypt_password(encrypted: &[u8]) -> Option<String> {
+    if encrypted.len() < 12 {
+        return None;
+    }
+    let key = derive_key();
+    let cipher = Aes256Gcm::new(&key);
+    let nonce = Nonce::from_slice(&encrypted[..12]);
+    cipher.decrypt(nonce, &encrypted[12..]).ok().and_then(|plaintext| String::from_utf8(plaintext).ok())
+}
diff --git a/src/bin/client/repl.rs b/src/bin/client/repl.rs
index 1512091..5559cfd 100644
--- a/src/bin/client/repl.rs
+++ b/src/bin/client/repl.rs
@@ -2,10 +2,18 @@ use std::io::{BufRead, BufReader, Write, stdout};
 use std::net::TcpStream;
 use rusqlite::Connection;
 use std::time::Instant;
+use bcrypt::hash_with_salt;
 
 use crate::connection;
 use crate::database;
 
+const SALT: [u8; 16] = *b"fixedsalt1234567";
+
+fn hash_password(password: &str) -> String {
+    let result = hash_with_salt(password, 10, SALT).unwrap();
+    result.to_string()
+}
+
 use crossterm::{event::{read, Event, KeyCode, KeyModifiers}, terminal::{disable_raw_mode, enable_raw_mode, Clear, ClearType}, cursor::{MoveToColumn, Show}, execute};
 
 macro_rules! println {
@@ -47,7 +55,7 @@ impl ClientState {
 
     pub fn connect_to(&mut self, addr: &str) {
         self.addr = Some(addr.to_string());
-        match connection::connect(addr) {
+        match connection::connect(addr, &self.conn) {
             Ok(s) => {
                 self.stream = Some(s);
                 database::save_host(&self.conn, addr);
@@ -166,6 +174,108 @@ impl ClientState {
                                         }
                                     }
                                 }
+                                "register" => {
+                                    if self.stream.is_none() {
+                                        println!("\nNot connected. Use 'connect <ip>:<port>' first.");
+                                        continue 'outer;
+                                    }
+                                    let username = read_line("\nUsername: ");
+                                    if username.is_empty() {
+                                        continue 'outer;
+                                    }
+                                    let password = read_line("Password: ");
+                                    if password.is_empty() {
+                                        continue 'outer;
+                                    }
+                                    let hashed = hash_password(&password);
+                                    let msg = format!("REGISTER {} {}", username, hashed);
+                                    let s = self.stream.as_mut().unwrap();
+                                    let writer_stream = s.try_clone().unwrap();
+                                    let mut writer = writer_stream;
+                                    let response_reader_stream = s.try_clone().unwrap();
+                                    let mut response_reader = BufReader::new(response_reader_stream);
+                                    if let Err(_) = writeln!(writer, "{}", msg) {
+                                        self.stream = None;
+                                        println!("\nConnection lost");
+                                        continue 'outer;
+                                    }
+                                    let mut response = String::new();
+                                    if let Err(_) = response_reader.read_line(&mut response) {
+                                        self.stream = None;
+                                        println!("\nConnection lost");
+                                    } else {
+                                        println!("\n{}", response.trim());
+                                    }
+                                }
+                                "login" => {
+                                    if self.stream.is_none() {
+                                        println!("\nNot connected. Use 'connect <ip>:<port>' first.");
+                                        continue 'outer;
+                                    }
+                                    loop {
+                                        let username = read_line("\nUsername: ");
+                                        if username.is_empty() {
+                                            break;
+                                        }
+                                        let password = read_line("Password: ");
+                                        if password.is_empty() {
+                                            break;
+                                        }
+                                        let hashed = hash_password(&password);
+                                        let msg = format!("AUTH {} {}", username, hashed);
+                                        let s = self.stream.as_mut().unwrap();
+                                        let writer_stream = s.try_clone().unwrap();
+                                        let mut writer = writer_stream;
+                                        let response_reader_stream = s.try_clone().unwrap();
+                                        let mut response_reader = BufReader::new(response_reader_stream);
+                                        if let Err(_) = writeln!(writer, "{}", msg) {
+                                            self.stream = None;
+                                            println!("\nConnection lost");
+                                            break;
+                                        }
+                                        let mut response = String::new();
+                                        if let Err(_) = response_reader.read_line(&mut response) {
+                                            self.stream = None;
+                                            println!("\nConnection lost");
+                                            break;
+                                        } else {
+                                            let resp = response.trim();
+                                            println!("\n{}", resp);
+                                            if resp.starts_with("OK") {
+                                                // Save credentials
+                                                if let Some(addr) = &self.addr {
+                                                    if let Err(e) = database::save_credentials(&self.conn, addr, &username, &password) {
+                                                        println!("Failed to save credentials: {}", e);
+                                                    }
+                                                }
+                                                break;
+                                            } else {
+                                                println!("Authentication failed, try again or press Enter to cancel.");
+                                            }
+                                        }
+                                    }
+                                }
+                                "setcred" => {
+                                    if self.addr.is_none() {
+                                        println!("\nNo host connected. Use 'connect <ip>:<port>' first.");
+                                        continue 'outer;
+                                    }
+                                    let username = read_line("\nUsername: ");
+                                    if username.is_empty() {
+                                        continue 'outer;
+                                    }
+                                    let password = read_line("Password: ");
+                                    if password.is_empty() {
+                                        continue 'outer;
+                                    }
+                                    if let Some(addr) = &self.addr {
+                                        if let Err(e) = database::save_credentials(&self.conn, addr, &username, &password) {
+                                            println!("Failed to save credentials: {}", e);
+                                        } else {
+                                            println!("Credentials saved for {}", addr);
+                                        }
+                                    }
+                                }
                                 "help" => {
                                     help();
                                 }
@@ -208,6 +318,18 @@ fn redraw(current_input: &str) {
     execute!(stdout(), MoveToColumn((2 + current_input.len()) as u16)).unwrap();
 }
 
+fn read_line(prompt: &str) -> String {
+    disable_raw_mode().unwrap();
+    execute!(stdout(), Show).unwrap();
+    print!("{}", prompt);
+    stdout().flush().unwrap();
+    let mut input = String::new();
+    std::io::stdin().read_line(&mut input).unwrap();
+    enable_raw_mode().unwrap();
+    execute!(stdout(), crossterm::cursor::Hide).unwrap();
+    input.trim().to_string()
+}
+
 fn help() {
-    println!("Available commands: connect <ip>:<port>, send <message>, disconnect, ping, help, quit");
+    println!("Available commands: connect <ip>:<port>, send <message>, disconnect, ping, login, register, setcred, help, quit");
 }