From 00ea8dca71e5c1d1fe7bf3d51984bfecbf5c3352 Mon Sep 17 00:00:00 2001 From: Carson McManus Date: Fri, 18 Feb 2022 14:55:10 -0500 Subject: [PATCH] refactor: account manager no longer needs to have all accounts loaded in order to function --- src/accountmanager.rs | 375 ++++++++++++++++++++++++++++-------------- src/main.rs | 54 +++--- 2 files changed, 278 insertions(+), 151 deletions(-) diff --git a/src/accountmanager.rs b/src/accountmanager.rs index a9a0915..3f644b1 100644 --- a/src/accountmanager.rs +++ b/src/accountmanager.rs @@ -2,6 +2,7 @@ pub use crate::encryption::EntryEncryptionParams; use crate::encryption::EntryEncryptor; use log::*; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use std::fs::File; use std::io::{BufReader, Read, Write}; use std::path::Path; @@ -28,9 +29,11 @@ pub struct Manifest { pub auto_confirm_trades: bool, #[serde(skip)] - pub accounts: Vec>>, + accounts: HashMap>>, #[serde(skip)] folder: String, // I wanted to use a Path here, but it was too hard to make it work... + #[serde(skip)] + passkey: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -56,8 +59,9 @@ impl Default for Manifest { auto_confirm_market_transactions: false, auto_confirm_trades: false, - accounts: vec![], + accounts: HashMap::new(), folder: "".into(), + passkey: None, } } } @@ -80,39 +84,61 @@ impl Manifest { return Ok(manifest); } - pub fn load_accounts( - &mut self, - passkey: &Option, - ) -> anyhow::Result<(), ManifestAccountLoadError> { - for entry in &mut self.entries { - let path = Path::new(&self.folder).join(&entry.filename); - debug!("loading account: {:?}", path); - let file = File::open(path)?; - let mut reader = BufReader::new(file); - let account: SteamGuardAccount; - match (passkey, entry.encryption.as_ref()) { - (Some(passkey), Some(params)) => { - let mut ciphertext: Vec = vec![]; - reader.read_to_end(&mut ciphertext)?; - let plaintext = crate::encryption::LegacySdaCompatible::decrypt( - passkey, params, ciphertext, - )?; - if plaintext[0] != '{' as u8 && plaintext[plaintext.len() - 1] != '}' as u8 { - return Err(ManifestAccountLoadError::IncorrectPasskey); - } - let s = std::str::from_utf8(&plaintext).unwrap(); - account = serde_json::from_str(&s)?; - } - (None, Some(_)) => { - return Err(ManifestAccountLoadError::MissingPasskey); - } - (_, None) => { - account = serde_json::from_reader(reader)?; - } - }; - entry.account_name = account.account_name.clone(); - self.accounts.push(Arc::new(Mutex::new(account))); + /// Tells the manifest to keep track of the encryption passkey, and use it for encryption when loading or saving accounts. + pub fn submit_passkey(&mut self, passkey: Option) { + if passkey.is_some() { + debug!("passkey was submitted to manifest"); + } else { + debug!("passkey was revoked from manifest"); } + self.passkey = passkey; + } + + pub fn load_accounts(&mut self) -> anyhow::Result<(), ManifestAccountLoadError> { + let account_names: Vec = self + .entries + .iter() + .map(|entry| entry.account_name.clone()) + .collect(); + for account_name in account_names { + self.load_account(&account_name)?; + } + Ok(()) + } + + fn load_account( + &mut self, + account_name: &String, + ) -> anyhow::Result<(), ManifestAccountLoadError> { + let mut entry = self.get_entry_mut(account_name)?.clone(); + let path = Path::new(&self.folder).join(&entry.filename); + debug!("loading account: {:?}", path); + let file = File::open(path)?; + let mut reader = BufReader::new(file); + let account: SteamGuardAccount; + match (&self.passkey, entry.encryption.as_ref()) { + (Some(passkey), Some(params)) => { + let mut ciphertext: Vec = vec![]; + reader.read_to_end(&mut ciphertext)?; + let plaintext = + crate::encryption::LegacySdaCompatible::decrypt(&passkey, params, ciphertext)?; + if plaintext[0] != '{' as u8 && plaintext[plaintext.len() - 1] != '}' as u8 { + return Err(ManifestAccountLoadError::IncorrectPasskey); + } + let s = std::str::from_utf8(&plaintext).unwrap(); + account = serde_json::from_str(&s)?; + } + (None, Some(_)) => { + return Err(ManifestAccountLoadError::MissingPasskey); + } + (_, None) => { + account = serde_json::from_reader(reader)?; + } + }; + entry.account_name = account.account_name.clone(); + self.accounts + .insert(entry.account_name.clone(), Arc::new(Mutex::new(account))); + *self.get_entry_mut(account_name)? = entry; Ok(()) } @@ -134,7 +160,8 @@ impl Manifest { account_name: account.account_name.clone(), encryption: None, }); - self.accounts.push(Arc::new(Mutex::new(account))); + self.accounts + .insert(account.account_name.clone(), Arc::new(Mutex::new(account))); } pub fn import_account(&mut self, import_path: String) -> anyhow::Result<()> { @@ -156,33 +183,36 @@ impl Manifest { pub fn remove_account(&mut self, account_name: String) { let index = self - .accounts + .entries .iter() - .position(|a| a.lock().unwrap().account_name == account_name) + .position(|a| a.account_name == account_name) .unwrap(); - self.accounts.remove(index); + self.accounts.remove(&account_name); self.entries.remove(index); } - pub fn save(&self, passkey: &Option) -> anyhow::Result<()> { - ensure!( - self.entries.len() == self.accounts.len(), - "Manifest entries don't match accounts." - ); + /// Saves the manifest and all loaded accounts. + pub fn save(&self) -> anyhow::Result<()> { info!("Saving manifest and accounts..."); - for (entry, account) in self.entries.iter().zip(&self.accounts) { + for account in self + .accounts + .values() + .into_iter() + .map(|a| a.clone().lock().unwrap().clone()) + { + let entry = self.get_entry(&account.account_name)?.clone(); debug!("saving {}", entry.filename); - let serialized = serde_json::to_vec(account.as_ref())?; + let serialized = serde_json::to_vec(&account)?; ensure!( serialized.len() > 2, "Something extra weird happened and the account was serialized into nothing." ); let final_buffer: Vec; - match (passkey, entry.encryption.as_ref()) { + match (&self.passkey, entry.encryption.as_ref()) { (Some(passkey), Some(params)) => { final_buffer = crate::encryption::LegacySdaCompatible::encrypt( - passkey, params, serialized, + &passkey, params, serialized, )?; } (None, Some(_)) => { @@ -206,10 +236,68 @@ impl Manifest { file.sync_data()?; Ok(()) } + + /// Return all loaded accounts. Order is not guarenteed. + pub fn get_all_loaded(&self) -> Vec>> { + return self.accounts.values().cloned().into_iter().collect(); + } + + pub fn get_entry( + &self, + account_name: &String, + ) -> anyhow::Result<&ManifestEntry, ManifestAccountLoadError> { + self.entries + .iter() + .find(|e| &e.account_name == account_name) + .ok_or(ManifestAccountLoadError::MissingManifestEntry) + } + + pub fn get_entry_mut( + &mut self, + account_name: &String, + ) -> anyhow::Result<&mut ManifestEntry, ManifestAccountLoadError> { + self.entries + .iter_mut() + .find(|e| &e.account_name == account_name) + .ok_or(ManifestAccountLoadError::MissingManifestEntry) + } + + pub fn has_passkey(&self) -> bool { + self.passkey.is_some() + } + + /// Gets the specified account by name. + /// Fails if the account does not exist in the manifest entries. + pub fn get_account( + &self, + account_name: &String, + ) -> anyhow::Result>> { + let account = self + .accounts + .get(account_name) + .map(|a| a.clone()) + .ok_or(anyhow!("Account not loaded")); + return account; + } + + /// Get or load the spcified account. + pub fn get_or_load_account( + &mut self, + account_name: &String, + ) -> anyhow::Result>, ManifestAccountLoadError> { + let account = self.get_account(account_name); + if account.is_ok() { + return Ok(account.unwrap()); + } + self.load_account(&account_name)?; + return Ok(self.get_account(account_name)?); + } } #[derive(Debug, Error)] pub enum ManifestAccountLoadError { + #[error("Could not find an entry in the manifest for this account. Check your spelling.")] + MissingManifestEntry, #[error("Manifest accounts are encrypted, but no passkey was provided.")] MissingPasskey, #[error("Incorrect passkey provided.")] @@ -253,7 +341,7 @@ mod tests { let tmp_dir = TempDir::new("steamguard-cli-test").unwrap(); let manifest_path = tmp_dir.path().join("manifest.json"); let manifest = Manifest::new(manifest_path.as_path()); - assert!(matches!(manifest.save(&None), Ok(_))); + assert!(matches!(manifest.save(), Ok(_))); } #[test] @@ -269,26 +357,39 @@ mod tests { "zvIayp3JPvtvX/QGHqsqKBk/44s=".into(), )?; manifest.add_account(account); - manifest.save(&None)?; + manifest.save()?; let mut loaded_manifest = Manifest::load(manifest_path.as_path())?; assert_eq!(loaded_manifest.entries.len(), 1); assert_eq!(loaded_manifest.entries[0].filename, "asdf1234.maFile"); - loaded_manifest.load_accounts(&None)?; + loaded_manifest.load_accounts()?; assert_eq!( loaded_manifest.entries.len(), loaded_manifest.accounts.len() ); + let account_name = "asdf1234".into(); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().account_name, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .account_name, "asdf1234" ); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().revocation_code, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .revocation_code, "R12345" ); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().shared_secret, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .shared_secret, steamguard::token::TwoFactorSecret::parse_shared_secret( "zvIayp3JPvtvX/QGHqsqKBk/44s=".into() )?, @@ -297,9 +398,9 @@ mod tests { } #[test] - fn test_should_save_and_load_manifest_encrypted() { - let passkey: Option = Some("password".into()); - let tmp_dir = TempDir::new("steamguard-cli-test").unwrap(); + fn test_should_save_and_load_manifest_encrypted() -> anyhow::Result<()> { + let passkey = Some("password".into()); + let tmp_dir = TempDir::new("steamguard-cli-test")?; let manifest_path = tmp_dir.path().join("manifest.json"); let mut manifest = Manifest::new(manifest_path.as_path()); let mut account = SteamGuardAccount::new(); @@ -307,41 +408,60 @@ mod tests { account.revocation_code = "R12345".into(); account.shared_secret = steamguard::token::TwoFactorSecret::parse_shared_secret( "zvIayp3JPvtvX/QGHqsqKBk/44s=".into(), - ) - .unwrap(); + )?; manifest.add_account(account); manifest.entries[0].encryption = Some(EntryEncryptionParams::generate()); - assert!(matches!(manifest.save(&passkey), Ok(_))); + manifest.submit_passkey(passkey.clone()); + assert!(matches!(manifest.save(), Ok(_))); let mut loaded_manifest = Manifest::load(manifest_path.as_path()).unwrap(); + loaded_manifest.submit_passkey(passkey); assert_eq!(loaded_manifest.entries.len(), 1); assert_eq!(loaded_manifest.entries[0].filename, "asdf1234.maFile"); - assert!(matches!(loaded_manifest.load_accounts(&passkey), Ok(_))); + let _r = loaded_manifest.load_accounts(); + if _r.is_err() { + eprintln!("{:?}", _r); + } + assert!(matches!(_r, Ok(_))); assert_eq!( loaded_manifest.entries.len(), loaded_manifest.accounts.len() ); + let account_name = "asdf1234".into(); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().account_name, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .account_name, "asdf1234" ); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().revocation_code, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .revocation_code, "R12345" ); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().shared_secret, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .shared_secret, steamguard::token::TwoFactorSecret::parse_shared_secret( "zvIayp3JPvtvX/QGHqsqKBk/44s=".into() ) .unwrap(), ); + Ok(()) } #[test] fn test_should_save_and_load_manifest_encrypted_longer() -> anyhow::Result<()> { - let passkey: Option = Some("password".into()); - let tmp_dir = TempDir::new("steamguard-cli-test").unwrap(); + let passkey = Some("password".into()); + let tmp_dir = TempDir::new("steamguard-cli-test")?; let manifest_path = tmp_dir.path().join("manifest.json"); let mut manifest = Manifest::new(manifest_path.as_path()); let mut account = SteamGuardAccount::new(); @@ -354,27 +474,42 @@ mod tests { account.uri = "otpauth://;laksdjf;lkasdjf;lkasdj;flkasdjlkf;asjdlkfjslk;adjfl;kasdjf;lksdjflk;asjd;lfajs;ldkfjaslk;djf;lsakdjf;lksdj".into(); account.token_gid = "asdf1234".into(); manifest.add_account(account); + manifest.submit_passkey(passkey.clone()); manifest.entries[0].encryption = Some(EntryEncryptionParams::generate()); - manifest.save(&passkey)?; + manifest.save()?; let mut loaded_manifest = Manifest::load(manifest_path.as_path())?; + loaded_manifest.submit_passkey(passkey.clone()); assert_eq!(loaded_manifest.entries.len(), 1); assert_eq!(loaded_manifest.entries[0].filename, "asdf1234.maFile"); - loaded_manifest.load_accounts(&passkey)?; + loaded_manifest.load_accounts()?; assert_eq!( loaded_manifest.entries.len(), loaded_manifest.accounts.len() ); + let account_name = "asdf1234".into(); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().account_name, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .account_name, "asdf1234" ); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().revocation_code, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .revocation_code, "R12345" ); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().shared_secret, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .shared_secret, steamguard::token::TwoFactorSecret::parse_shared_secret( "zvIayp3JPvtvX/QGHqsqKBk/44s=".into() ) @@ -385,8 +520,8 @@ mod tests { } #[test] - fn test_should_import() { - let tmp_dir = TempDir::new("steamguard-cli-test").unwrap(); + fn test_should_import() -> anyhow::Result<()> { + let tmp_dir = TempDir::new("steamguard-cli-test")?; let manifest_path = tmp_dir.path().join("manifest.json"); let mut manifest = Manifest::new(manifest_path.as_path()); let mut account = SteamGuardAccount::new(); @@ -397,8 +532,8 @@ mod tests { ) .unwrap(); manifest.add_account(account); - assert!(matches!(manifest.save(&None), Ok(_))); - std::fs::remove_file(&manifest_path).unwrap(); + manifest.save()?; + std::fs::remove_file(&manifest_path)?; let mut loaded_manifest = Manifest::new(manifest_path.as_path()); assert!(matches!( @@ -416,104 +551,94 @@ mod tests { loaded_manifest.entries.len(), loaded_manifest.accounts.len() ); + let account_name = "asdf1234".into(); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().account_name, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .account_name, "asdf1234" ); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().revocation_code, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .revocation_code, "R12345" ); assert_eq!( - loaded_manifest.accounts[0].lock().unwrap().shared_secret, + loaded_manifest + .get_account(&account_name)? + .lock() + .unwrap() + .shared_secret, steamguard::token::TwoFactorSecret::parse_shared_secret( "zvIayp3JPvtvX/QGHqsqKBk/44s=".into() ) .unwrap(), ); + + Ok(()) } #[test] - fn test_sda_compatibility_1() { + fn test_sda_compatibility_1() -> anyhow::Result<()> { let path = Path::new("src/fixtures/maFiles/compat/1-account/manifest.json"); assert!(path.is_file()); - let result = Manifest::load(path); - assert!(matches!(result, Ok(_))); - let mut manifest = result.unwrap(); + let mut manifest = Manifest::load(path)?; assert!(matches!(manifest.entries.last().unwrap().encryption, None)); - assert!(matches!(manifest.load_accounts(&None), Ok(_))); + manifest.load_accounts()?; + let account_name = manifest.entries.last().unwrap().account_name.clone(); assert_eq!( - manifest.entries.last().unwrap().account_name, + account_name, manifest - .accounts - .last() - .unwrap() + .get_account(&account_name)? .lock() .unwrap() .account_name ); + Ok(()) } #[test] - fn test_sda_compatibility_1_encrypted() { + fn test_sda_compatibility_1_encrypted() -> anyhow::Result<()> { let path = Path::new("src/fixtures/maFiles/compat/1-account-encrypted/manifest.json"); assert!(path.is_file()); - let result = Manifest::load(path); - assert!(matches!(result, Ok(_))); - let mut manifest = result.unwrap(); + let mut manifest = Manifest::load(path)?; assert!(matches!( manifest.entries.last().unwrap().encryption, Some(_) )); - let result = manifest.load_accounts(&Some("password".into())); - assert!( - matches!(result, Ok(_)), - "error when loading accounts: {:?}", - result.unwrap_err() - ); + manifest.submit_passkey(Some("password".into())); + manifest.load_accounts()?; + let account_name = manifest.entries.last().unwrap().account_name.clone(); assert_eq!( - manifest.entries.last().unwrap().account_name, + account_name, manifest - .accounts - .last() - .unwrap() + .get_account(&account_name)? .lock() .unwrap() .account_name ); + Ok(()) } #[test] - fn test_sda_compatibility_no_webcookie() { + fn test_sda_compatibility_no_webcookie() -> anyhow::Result<()> { let path = Path::new("src/fixtures/maFiles/compat/no-webcookie/manifest.json"); assert!(path.is_file()); - let result = Manifest::load(path); - assert!(matches!(result, Ok(_))); - let mut manifest = result.unwrap(); + let mut manifest = Manifest::load(path)?; assert!(matches!(manifest.entries.last().unwrap().encryption, None)); - assert!(matches!(manifest.load_accounts(&None), Ok(_))); + assert!(matches!(manifest.load_accounts(), Ok(_))); + let account_name = manifest.entries.last().unwrap().account_name.clone(); + let account = manifest.get_account(&account_name)?; + assert_eq!(account_name, account.lock().unwrap().account_name); assert_eq!( - manifest.entries.last().unwrap().account_name, - manifest - .accounts - .last() - .unwrap() - .lock() - .unwrap() - .account_name - ); - assert_eq!( - manifest - .accounts - .last() - .unwrap() - .lock() - .unwrap() - .session - .as_ref() - .unwrap() - .web_cookie, + account.lock().unwrap().session.as_ref().unwrap().web_cookie, None ); + Ok(()) } } diff --git a/src/main.rs b/src/main.rs index d064b06..0e3138e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -190,24 +190,26 @@ fn run() -> anyhow::Result<()> { std::fs::create_dir_all(mafiles_dir)?; manifest = accountmanager::Manifest::new(path.as_path()); - manifest.save(&None)?; + manifest.save()?; } else { manifest = accountmanager::Manifest::load(path.as_path())?; } let mut passkey: Option = matches.value_of("passkey").map(|s| s.into()); + manifest.submit_passkey(passkey); loop { - match manifest.load_accounts(&passkey) { + match manifest.load_accounts() { Ok(_) => break, Err( accountmanager::ManifestAccountLoadError::MissingPasskey | accountmanager::ManifestAccountLoadError::IncorrectPasskey, ) => { - if passkey.is_some() { + if manifest.has_passkey() { error!("Incorrect passkey"); } passkey = rpassword::prompt_password_stdout("Enter encryption passkey: ").ok(); + manifest.submit_passkey(passkey); } Err(e) => { error!("Could not load accounts: {}", e); @@ -220,6 +222,7 @@ fn run() -> anyhow::Result<()> { println!("Log in to the account that you want to link to steamguard-cli"); print!("Username: "); let username = tui::prompt(); + let account_name = username.clone(); if manifest.account_exists(&username) { bail!( "Account {} already exists in manifest, remove it first", @@ -264,26 +267,20 @@ fn run() -> anyhow::Result<()> { } } manifest.add_account(account); - match manifest.save(&passkey) { + match manifest.save() { Ok(_) => {} Err(err) => { error!("Aborting the account linking process because we failed to save the manifest. This is really bad. Here is the error: {}", err); println!( "Just in case, here is the account info. Save it somewhere just in case!\n{:?}", - manifest.accounts.last().unwrap().lock().unwrap() + manifest.get_account(&account_name).unwrap().lock().unwrap() ); return Err(err.into()); } } - let mut account = manifest - .accounts - .last() - .as_ref() - .unwrap() - .clone() - .lock() - .unwrap(); + let account_arc = manifest.get_account(&account_name).unwrap(); + let mut account = account_arc.lock().unwrap(); println!("Authenticator has not yet been linked. Before continuing with finalization, please take the time to write down your revocation code: {}", account.revocation_code); tui::pause(); @@ -311,7 +308,7 @@ fn run() -> anyhow::Result<()> { } println!("Authenticator finalized."); - match manifest.save(&None) { + match manifest.save() { Ok(_) => {} Err(err) => { println!( @@ -340,10 +337,10 @@ fn run() -> anyhow::Result<()> { } } - manifest.save(&passkey)?; + manifest.save()?; return Ok(()); } else if matches.is_present("encrypt") { - if passkey.is_none() { + if !manifest.has_passkey() { loop { passkey = rpassword::prompt_password_stdout("Enter encryption passkey: ").ok(); let passkey_confirm = @@ -353,34 +350,39 @@ fn run() -> anyhow::Result<()> { } error!("Passkeys do not match, try again."); } + manifest.submit_passkey(passkey); } for entry in &mut manifest.entries { entry.encryption = Some(accountmanager::EntryEncryptionParams::generate()); } - manifest.save(&passkey)?; + manifest.save()?; return Ok(()); } else if matches.is_present("decrypt") { for entry in &mut manifest.entries { entry.encryption = None; } - manifest.save(&passkey)?; + manifest.submit_passkey(None); + manifest.save()?; return Ok(()); } let mut selected_accounts: Vec>> = vec![]; if matches.is_present("all") { + manifest + .load_accounts() + .expect("Failed to load all requested accounts, aborting"); // manifest.accounts.iter().map(|a| selected_accounts.push(a.b)); - for account in &manifest.accounts { - selected_accounts.push(account.clone()); + for entry in &manifest.entries { + selected_accounts.push(manifest.get_account(&entry.account_name).unwrap().clone()); } } else { - for account in &manifest.accounts { + for entry in &manifest.entries { if !matches.is_present("username") { - selected_accounts.push(account.clone()); + selected_accounts.push(manifest.get_account(&entry.account_name).unwrap().clone()); break; } - if matches.value_of("username").unwrap() == account.lock().unwrap().account_name { - selected_accounts.push(account.clone()); + if matches.value_of("username").unwrap() == entry.account_name { + selected_accounts.push(manifest.get_account(&entry.account_name).unwrap().clone()); break; } } @@ -471,7 +473,7 @@ fn run() -> anyhow::Result<()> { } } - manifest.save(&passkey)?; + manifest.save()?; } else if let Some(_) = matches.subcommand_matches("remove") { println!( "This will remove the mobile authenticator from {} accounts: {}", @@ -528,7 +530,7 @@ fn run() -> anyhow::Result<()> { manifest.remove_account(account_name); } - manifest.save(&passkey)?; + manifest.save()?; } else { let server_time = steamapi::get_server_time(); debug!("Time used to generate codes: {}", server_time);