From 85dc841823813c1d4caa2bfb608fa7a05ab98a2c Mon Sep 17 00:00:00 2001 From: "Christopher L. Crutchfield" Date: Mon, 18 Nov 2024 15:01:57 -0800 Subject: [PATCH] feat: unit tests for credential manager --- src/credential_manager.rs | 260 +++++++++++++++++++++------ src/subcommands/login_subcommand.rs | 2 +- src/subcommands/logout_subcommand.rs | 2 +- 3 files changed, 204 insertions(+), 60 deletions(-) diff --git a/src/credential_manager.rs b/src/credential_manager.rs index cf492a5..7bcf621 100644 --- a/src/credential_manager.rs +++ b/src/credential_manager.rs @@ -1,14 +1,13 @@ -use std::{fs::create_dir_all, process::Command}; +use std::{collections::HashMap, fs::create_dir_all, process::Command}; use aes_gcm::{aead::{Aead, OsRng}, AeadCore, Aes256Gcm, Key, KeyInit, Nonce}; use app_dirs2::{AppDataType, AppInfo, app_root}; use anyhow::{Result, Context}; -use keyring::Entry; +use keyring::{credential, Entry}; use rusqlite::Connection; #[derive(Debug)] struct DatabaseCredential { - id: i32, url: String, user: String, totp_comand_encrypted: Option>, @@ -22,10 +21,19 @@ pub struct Credential { } impl Credential { - pub fn new(user: &str, password: &str, totp_command: Option) -> Credential { + pub fn new(user: &str, password: &str, totp_command: Option<&str>) -> Credential { + let totp_command = match totp_command { + Some(totp_command) => Some(totp_command.to_string()), + None => None + }; + + Credential::new_with_string(user.to_string(), password.to_string(), totp_command) + } + + pub fn new_with_string(user: String, password: String, totp_command: Option) -> Credential { Credential { - user: user.to_string(), - password: password.to_string(), + user, + password, totp_command } } @@ -50,14 +58,19 @@ impl Credential { } pub struct CredentialManager { + connection: Connection, + entry_cache: HashMap<(String, String), Entry> } impl CredentialManager { - pub fn new() -> CredentialManager { - CredentialManager { } + pub fn new() -> Result { + Ok(CredentialManager { + connection: CredentialManager::get_connection()?, + entry_cache: HashMap::new() + }) } - fn get_database(&self) -> Result { + fn get_connection() -> Result { // Get the path to the credential database let mut path = app_root(AppDataType::UserConfig, &AppInfo{ name: "git-lfs-synology", @@ -71,59 +84,76 @@ impl CredentialManager { create_dir_all(sqlite_path.parent().context("No parent")?)?; } - let should_init_database = !sqlite_path.exists(); - let conn = Connection::open(sqlite_path)?; - - if should_init_database { - println!("Creating table"); - conn.execute( - "CREATE TABLE IF NOT EXISTS Credentials ( - id INTEGER PRIMARY KEY, - url TEXT NOT NULL, - user TEXT NOT NULL, - totp_command_encrypted BLOB, - totp_nonce BLOB - )", - (), // empty list of parameters. - )?; - } - - Ok(conn) + Ok(Connection::open(sqlite_path)?) } fn get_database_credential_iter(&self, url: &str) -> Result> { let database = self.get_database()?; let mut stmt: rusqlite::Statement<'_> = database.prepare( - "SELECT id, url, user, totp_command_encrypted, totp_nonce FROM Credentials WHERE url=:url;")?; + "SELECT url, user, totp_command_encrypted, totp_nonce FROM Credentials WHERE url=:url;")?; let rows: Vec = stmt.query_map(&[(":url", url)], |row| { Ok(DatabaseCredential { - id: row.get(0)?, - url: row.get(1)?, - user: row.get(2)?, - totp_comand_encrypted: row.get(3)?, - totp_nonce: row.get(4)? + url: row.get(0)?, + user: row.get(1)?, + totp_comand_encrypted: row.get(2)?, + totp_nonce: row.get(3)? }) })?.filter_map(|r| r.ok()).collect::>().try_into()?; + + Ok(rows) + } - drop(stmt); // Allow closing the database. + fn get_database(&self) -> Result<&Connection> { + let conn = &self.connection; + conn.execute( + "CREATE TABLE IF NOT EXISTS Credentials ( + id INTEGER PRIMARY KEY, + url TEXT NOT NULL, + user TEXT NOT NULL, + totp_command_encrypted BLOB, + totp_nonce BLOB + )", + (), // empty list of parameters. + )?; - match database.close() { - Ok(_) => Ok(rows), - Err(_) => Err(anyhow::Error::msg("An error occurred closig the database.")) + Ok(conn) + } + + fn get_entry(&mut self, url: &str, user: &str) -> Result<&Entry>{ + if !self.entry_cache.contains_key(&(url.to_string(), user.to_string())) { + let entry = Entry::new(url, user)?; + self.entry_cache.insert((url.to_string(), user.to_string()), entry); } + + Ok(self.entry_cache.get(&(url.to_string(), user.to_string())).context("Entry does not exist in cache")?) } - pub fn get_credential(&self, url: &str) -> Result { + fn pad_string(&self, input: &str) -> String { + let mut output = input.to_string(); + + while output.len() < 32 { + output.push(' '); + } + + output + } + + pub fn get_credential(&mut self, url: &str) -> Result> { + if !self.has_credential(url)? { + return Ok(None); + } + let database_rows = self.get_database_credential_iter(url)?; let database_credential = database_rows.first().context("No elements returned from database.")?; - let entry = Entry::new(url, &database_credential.user)?; + let entry = self.get_entry(url, &database_credential.user)?; let password = entry.get_password()?; - let mut totp_command: Option = None; + let mut totp_command: Option = None; if database_credential.totp_comand_encrypted.is_some() { - let key: &Key = password.as_bytes().into(); + let padded_password = self.pad_string(password.as_str()); + let key: &Key = padded_password.as_bytes().into(); let nonce_vec = database_credential.totp_nonce.clone().context("No nonce provided for credential")?; @@ -135,7 +165,7 @@ impl CredentialManager { totp_command = Some(String::from_utf8(plaintext)?); } - Ok(Credential::new(&database_credential.user, &password, totp_command)) + Ok(Some(Credential::new_with_string(database_credential.user.clone(), password, totp_command))) } pub fn has_credential(&self, url: &str) -> Result { @@ -144,13 +174,14 @@ impl CredentialManager { Ok(!database_rows.is_empty()) } - pub fn remove_credential(&self, url: &str) -> Result<()> { + pub fn remove_credential(&mut self, url: &str) -> Result<()> { if self.has_credential(url)? { let database_rows = self.get_database_credential_iter(url)?; let database_credential = database_rows.first().context("No elements returned from database.")?; - let entry = Entry::new(&database_credential.url, &database_credential.user)?; + let entry = self.get_entry(url, &database_credential.user)?; entry.delete_credential()?; + self.entry_cache.remove(&(url.to_string(), database_credential.user.to_string())); let database = self.get_database()?; @@ -163,7 +194,7 @@ impl CredentialManager { Ok(()) } - pub fn set_credential(&self, url: &str, credential: &Credential) -> Result<()> { + pub fn set_credential(&mut self, url: &str, credential: &Credential) -> Result<()> { if self.has_credential(url)? { self.remove_credential(url)?; } @@ -171,7 +202,8 @@ impl CredentialManager { let mut totp_comand_encrypted: Option> = None; let mut totp_nonce: Option> = None; if credential.totp_command.is_some() { - let key: &Key = credential.password.as_bytes().into(); + let padded_password = self.pad_string(credential.password.as_str()); + let key: &Key = padded_password.as_bytes().into(); let cipher = Aes256Gcm::new(&key); let nonce = Aes256Gcm::generate_nonce(&mut OsRng); // 96-bits; unique per message @@ -183,26 +215,138 @@ impl CredentialManager { let database = self.get_database()?; database.execute( - "INSERT INTO Credential (url, user, totp_command_encrypted, totp_nonce) VALUES (?1)", - (url.to_string(), credential.user.to_string(), totp_comand_encrypted, totp_nonce), - )?; - - let entry = Entry::new(url, &credential.user)?; + "INSERT INTO Credentials (url, user, totp_command_encrypted, totp_nonce) VALUES (?1, ?2, ?3, ?4)", + ( + url.to_string(), + credential.user.to_string(), + totp_comand_encrypted, + totp_nonce, + ))?; + + let entry = self.get_entry(url, &credential.user)?; entry.set_password(&credential.password)?; - - match database.close() { - Ok(_) => Ok(()), - Err(_) => Err(anyhow::Error::msg("An error occurred closig the database.")) - } + + Ok(()) } } #[cfg(test)] mod tests { - use super::CredentialManager; + use std::collections::HashMap; + + use anyhow::Context; + use keyring::{mock, set_default_credential_builder}; + use rusqlite::Connection; + + use super::{CredentialManager, Credential}; #[test] fn create_credential_manager() { - let credential_manager = CredentialManager::new(); + let _ = CredentialManager::new(); + } + + #[test] + fn set_get_credential_no_totp_command() { + set_default_credential_builder(mock::default_credential_builder()); // Set mock + + let mut credential_manager = CredentialManager { + connection: Connection::open_in_memory().unwrap(), + entry_cache: HashMap::new() + }; + + credential_manager.set_credential("http://example.com", &Credential::new("test_user", "test_password", None)).unwrap(); + + let credential: Credential = credential_manager.get_credential("http://example.com").unwrap().context("Credential expected").unwrap(); + + assert_eq!(credential.user, "test_user".to_string()); + assert_eq!(credential.password, "test_password".to_string()); + assert_eq!(credential.totp_command, None); + } + + #[test] + fn set_get_credential_totp_command() { + set_default_credential_builder(mock::default_credential_builder()); // Set mock + + let mut credential_manager = CredentialManager { + connection: Connection::open_in_memory().unwrap(), + entry_cache: HashMap::new() + }; + + credential_manager.set_credential("http://example.com", &Credential::new("test_user", "test_password", Some("echo 12345"))).unwrap(); + + let credential: Credential = credential_manager.get_credential("http://example.com").unwrap().context("Credential expected").unwrap(); + + assert_eq!(credential.user, "test_user".to_string()); + assert_eq!(credential.password, "test_password".to_string()); + assert_eq!(credential.totp_command.context("Should not be null").unwrap(), "echo 12345".to_string()); + } + + #[test] + fn has_credential() { + set_default_credential_builder(mock::default_credential_builder()); // Set mock + + let mut credential_manager = CredentialManager { + connection: Connection::open_in_memory().unwrap(), + entry_cache: HashMap::new() + }; + + credential_manager.set_credential("http://example.com", &Credential::new("test_user", "test_password", Some("echo 12345"))).unwrap(); + + assert!(credential_manager.has_credential("http://example.com").unwrap()); + } + + #[test] + fn can_update_credential_no_totp() { + set_default_credential_builder(mock::default_credential_builder()); // Set mock + + let mut credential_manager = CredentialManager { + connection: Connection::open_in_memory().unwrap(), + entry_cache: HashMap::new() + }; + + credential_manager.set_credential("http://example.com", &Credential::new("test_user", "test_password", Some("echo 12345"))).unwrap(); + + credential_manager.set_credential("http://example.com", &Credential::new("test_user", "test_password2", None)).unwrap(); + + let credential: Credential = credential_manager.get_credential("http://example.com").unwrap().context("Credential expected").unwrap(); + + assert_eq!(credential.user, "test_user".to_string()); + assert_eq!(credential.password, "test_password2".to_string()); + assert_eq!(credential.totp_command, None); + } + + #[test] + fn can_update_credential_totp() { + set_default_credential_builder(mock::default_credential_builder()); // Set mock + + let mut credential_manager = CredentialManager { + connection: Connection::open_in_memory().unwrap(), + entry_cache: HashMap::new() + }; + + credential_manager.set_credential("http://example.com", &Credential::new("test_user", "test_password", Some("echo 12345"))).unwrap(); + + credential_manager.set_credential("http://example.com", &Credential::new("test_user", "test_password2", Some("echo 123456"))).unwrap(); + + let credential: Credential = credential_manager.get_credential("http://example.com").unwrap().context("Credential expected").unwrap(); + + assert_eq!(credential.user, "test_user".to_string()); + assert_eq!(credential.password, "test_password2".to_string()); + assert_eq!(credential.totp_command.context("TOTP should of not null").unwrap(), "echo 123456".to_string()); + } + + #[test] + fn remove_credential() { + set_default_credential_builder(mock::default_credential_builder()); // Set mock + + let mut credential_manager = CredentialManager { + connection: Connection::open_in_memory().unwrap(), + entry_cache: HashMap::new() + }; + + credential_manager.set_credential("http://example.com", &Credential::new("test_user", "test_password", Some("echo 12345"))).unwrap(); + credential_manager.remove_credential("http://example.com").unwrap(); + + assert!(!credential_manager.has_credential("http://example.com").unwrap()); } } \ No newline at end of file diff --git a/src/subcommands/login_subcommand.rs b/src/subcommands/login_subcommand.rs index 5bd0319..836b4e2 100644 --- a/src/subcommands/login_subcommand.rs +++ b/src/subcommands/login_subcommand.rs @@ -12,7 +12,7 @@ impl Subcommand for LoginSubcommand { let url = arg_matches.get_one::("URL").context("URL not provided.")?; let user = arg_matches.get_one::("USER").context("USER not provided.")?; - let credential_manager = CredentialManager::new(); + let credential_manager = CredentialManager::new()?; if !credential_manager.has_credential(url)? { // TODO need to ask for password from user diff --git a/src/subcommands/logout_subcommand.rs b/src/subcommands/logout_subcommand.rs index 94846ef..0579fce 100644 --- a/src/subcommands/logout_subcommand.rs +++ b/src/subcommands/logout_subcommand.rs @@ -11,7 +11,7 @@ impl Subcommand for LogoutSubcommand { fn execute(&self, arg_matches: &ArgMatches) -> Result<()> { let url = arg_matches.get_one::("URL").context("URL not provided.")?; - let credential_manager = CredentialManager::new(); + let mut credential_manager = CredentialManager::new()?; credential_manager.remove_credential(url)?; Ok(())