1 files changed, 94 insertions, 0 deletions

diff --git a/src/client/database.rs b/src/client/database.rs
new file mode 100644
index 0000000..03a6c45
--- /dev/null
+++ b/src/client/database.rs
@@ -0,0 +1,94 @@
+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();
+    conn.execute(
+        "CREATE TABLE IF NOT EXISTS hosts (
+            host TEXT PRIMARY KEY,
+            last_used DATETIME DEFAULT CURRENT_TIMESTAMP
+        )",
+        [],
+    ).unwrap();
+    conn.execute(
+        "CREATE TABLE IF NOT EXISTS credentials (
+            host TEXT PRIMARY KEY,
+            username TEXT NOT NULL,
+            encrypted_password BLOB NOT NULL
+        )",
+        [],
+    ).unwrap();
+    conn
+}
+
+pub fn get_last_host(conn: &Connection) -> Option<String> {
+    let mut stmt = conn.prepare("SELECT host FROM hosts ORDER BY last_used DESC LIMIT 1").unwrap();
+    let mut rows = stmt.query([]).unwrap();
+    if let Ok(Some(row)) = rows.next() {
+        row.get(0).ok()
+    } else {
+        None
+    }
+}
+
+pub fn save_host(conn: &Connection, host: &str) {
+    conn.execute(
+        "INSERT OR REPLACE INTO hosts (host, last_used) VALUES (?1, CURRENT_TIMESTAMP)",
+        [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())
+}