7 files changed, 300 insertions, 7 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");
 }
diff --git a/src/bin/server/database.rs b/src/bin/server/database.rs
new file mode 100644
index 0000000..9c51997
--- /dev/null
+++ b/src/bin/server/database.rs
@@ -0,0 +1,31 @@
+use rusqlite::Connection;
+
+pub fn init_db() -> Connection {
+    let conn = Connection::open("server.db").unwrap();
+    conn.execute(
+        "CREATE TABLE IF NOT EXISTS users (
+            username TEXT PRIMARY KEY,
+            hashed_password TEXT NOT NULL
+        )",
+        [],
+    ).unwrap();
+    conn
+}
+
+pub fn add_user(conn: &Connection, username: &str, hashed_password: &str) -> Result<(), Box<dyn std::error::Error>> {
+    conn.execute(
+        "INSERT INTO users (username, hashed_password) VALUES (?1, ?2)",
+        [username, hashed_password],
+    )?;
+    Ok(())
+}
+
+pub fn get_hashed_password(conn: &Connection, username: &str) -> Option<String> {
+    let mut stmt = conn.prepare("SELECT hashed_password FROM users WHERE username = ?1").unwrap();
+    let mut rows = stmt.query([username]).unwrap();
+    if let Ok(Some(row)) = rows.next() {
+        row.get(0).ok()
+    } else {
+        None
+    }
+}
diff --git a/src/bin/server/handler.rs b/src/bin/server/handler.rs
index 906a91f..2132d09 100644
--- a/src/bin/server/handler.rs
+++ b/src/bin/server/handler.rs
@@ -1,15 +1,64 @@
 use std::io::{BufRead, BufReader, Write};
 use std::net::TcpStream;
+use std::sync::Mutex;
+use once_cell::sync::Lazy;
+use rusqlite::Connection;
+use crate::database;
+
+static DB_CONN: Lazy<Mutex<Connection>> = Lazy::new(|| Mutex::new(database::init_db()));
 
 pub fn handle_client(stream: TcpStream) {
-    println!("Client connected");
+    let client_addr = stream.peer_addr().unwrap().to_string();
+    let mut user = "guest".to_string();
+    println!("[{} @ {}] Client connected", user, client_addr);
     let reader_stream = stream.try_clone().unwrap();
     let reader = BufReader::new(reader_stream);
     let mut writer = stream;
     for line in reader.lines() {
         let line = line.unwrap();
-        println!("Received: {}", line);
-        writeln!(writer, "{}", line).unwrap();
+        println!("[{} @ {}] Received: {}", user, client_addr, line);
+        let response = process_message(&line, &mut user);
+        writeln!(writer, "{}", response).unwrap();
+    }
+    println!("[{} @ {}] Client disconnected", user, client_addr);
+}
+
+fn process_message(msg: &str, user: &mut String) -> String {
+    let parts: Vec<&str> = msg.splitn(3, ' ').collect();
+    match parts[0] {
+        "REGISTER" => {
+            if parts.len() != 3 {
+                "FAIL Invalid REGISTER format".to_string()
+            } else {
+                let username = parts[1];
+                let hashed_pass = parts[2];
+                let db = DB_CONN.lock().unwrap();
+                match database::add_user(&db, username, hashed_pass) {
+                    Ok(_) => format!("OK Registered {}", username),
+                    Err(_) => "FAIL User already exists".to_string(),
+                }
+            }
+        }
+        "AUTH" => {
+            if parts.len() != 3 {
+                "FAIL Invalid AUTH format".to_string()
+            } else {
+                let username = parts[1];
+                let provided_hash = parts[2];
+                let db = DB_CONN.lock().unwrap();
+                if let Some(stored_hash) = database::get_hashed_password(&db, username) {
+                    if provided_hash == stored_hash {
+                        *user = username.to_string();
+                        format!("OK {}", username)
+                    } else {
+                        "FAIL Invalid credentials".to_string()
+                    }
+                } else {
+                    "FAIL User not found".to_string()
+                }
+            }
+        }
+        "PING" => "PING".to_string(),
+        _ => msg.to_string(), // Echo back for compatibility
     }
-    println!("Client disconnected");
 }
diff --git a/src/bin/server/main.rs b/src/bin/server/main.rs
index 2c6e1dc..7693069 100644
--- a/src/bin/server/main.rs
+++ b/src/bin/server/main.rs
@@ -1,5 +1,6 @@
 mod args;
 mod handler;
+mod database;
 
 use std::net::TcpListener;
 use std::thread;
diff --git a/src/bin/server/mod.rs b/src/bin/server/mod.rs
index adb60bb..b2d622e 100644
--- a/src/bin/server/mod.rs
+++ b/src/bin/server/mod.rs
@@ -1,2 +1,3 @@
 pub mod args;
 pub mod handler;
+pub mod database;