Merge pull request #138 from dyc3/manifest-upgrades

Manifest auto-upgrades and Lazy account loading

src/accountmanager.rs

@@ -94,23 +94,36 @@ impl Manifest {
 		self.passkey = passkey;
 	}
 
+	/// Loads all accounts, registers them, and performs auto upgrades.
 	pub fn load_accounts(&mut self) -> anyhow::Result<(), ManifestAccountLoadError> {
-		let account_names: Vec<String> = self
-			.entries
-			.iter()
-			.map(|entry| entry.account_name.clone())
-			.collect();
-		for account_name in account_names {
-			self.load_account(&account_name)?;
+		self.auto_upgrade()?;
+		let mut accounts = vec![];
+		for entry in &self.entries {
+			let account = self.load_account_by_entry(&entry)?;
+			accounts.push(account);
+		}
+		for account in accounts {
+			self.register_loaded_account(account);
 		}
 		Ok(())
 	}
 
+	/// Loads an account by account name.
+	/// Must call `register_loaded_account` after loading the account.
 	fn load_account(
-		&mut self,
+		&self,
 		account_name: &String,
-	) -> anyhow::Result<(), ManifestAccountLoadError> {
-		let mut entry = self.get_entry_mut(account_name)?.clone();
+	) -> anyhow::Result<Arc<Mutex<SteamGuardAccount>>, ManifestAccountLoadError> {
+		let entry = self.get_entry(account_name)?;
+		self.load_account_by_entry(&entry)
+	}
+
+	/// Loads an account from a manifest entry.
+	/// Must call `register_loaded_account` after loading the account.
+	fn load_account_by_entry(
+		&self,
+		entry: &ManifestEntry,
+	) -> anyhow::Result<Arc<Mutex<SteamGuardAccount>>, ManifestAccountLoadError> {
 		let path = Path::new(&self.folder).join(&entry.filename);
 		debug!("loading account: {:?}", path);
 		let file = File::open(path)?;
@@ -135,11 +148,14 @@ impl Manifest {
 				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(())
+		let account = Arc::new(Mutex::new(account));
+		Ok(account)
+	}
+
+	/// Register an account as loaded, so it can be operated on.
+	fn register_loaded_account(&mut self, account: Arc<Mutex<SteamGuardAccount>>) {
+		let account_name = account.lock().unwrap().account_name.clone();
+		self.accounts.insert(account_name, account);
 	}
 
 	pub fn account_exists(&self, account_name: &String) -> bool {
@@ -289,8 +305,30 @@ impl Manifest {
 		if account.is_ok() {
 			return Ok(account.unwrap());
 		}
-		self.load_account(&account_name)?;
-		return Ok(self.get_account(account_name)?);
+		let account = self.load_account(&account_name)?;
+		self.register_loaded_account(account.clone());
+		return Ok(account);
+	}
+
+	/// Determine if any manifest entries are missing `account_name`.
+	fn is_missing_account_name(&self) -> bool {
+		self.entries.iter().any(|e| e.account_name.is_empty())
+	}
+
+	/// Performs auto-upgrades on the manifest. Returns true if any upgrades were performed.
+	pub fn auto_upgrade(&mut self) -> anyhow::Result<bool, ManifestAccountLoadError> {
+		debug!("Performing auto-upgrade...");
+		let mut upgraded = false;
+		if self.is_missing_account_name() {
+			debug!("Adding missing account names");
+			for i in 0..self.entries.len() {
+				let account = self.load_account_by_entry(&self.entries[i].clone())?;
+				self.entries[i].account_name = account.lock().unwrap().account_name.clone();
+			}
+			upgraded = true;
+		}
+
+		Ok(upgraded)
 	}
 }
 
@@ -668,4 +706,22 @@ mod tests {
 		);
 		Ok(())
 	}
+
+	#[cfg(test)]
+	mod manifest_upgrades {
+		use super::*;
+
+		#[test]
+		fn test_missing_account_name() {
+			let path = Path::new("src/fixtures/maFiles/compat/missing-account-name/manifest.json");
+			assert!(path.is_file());
+			let mut manifest = Manifest::load(path).unwrap();
+			assert_eq!(manifest.entries.len(), 1);
+			assert_eq!(manifest.entries[0].account_name, "".to_string());
+			assert!(manifest.is_missing_account_name());
+
+			manifest.auto_upgrade().unwrap();
+			assert_eq!(manifest.entries[0].account_name, "example".to_string());
+		}
+	}
 }

src/fixtures/maFiles/compat/missing-account-name/1234.maFile

@@ -0,0 +1 @@
+{"shared_secret":"zvIayp3JPvtvX/QGHqsqKBk/44s=","serial_number":"kljasfhds","revocation_code":"R12345","uri":"otpauth://totp/Steam:example?secret=ASDF&issuer=Steam","server_time":1602522478,"account_name":"example","token_gid":"jkkjlhkhjgf","identity_secret":"kjsdlwowiqe=","secret_1":"sklduhfgsdlkjhf=","status":1,"device_id":"android:99d2ad0e-4bad-4247-b111-26393aae0be3","fully_enrolled":true,"Session":{"SessionID":"a;lskdjf","SteamLogin":"983498437543","SteamLoginSecure":"dlkjdsl;j%7C%32984730298","WebCookie":";lkjsed;klfjas98093","OAuthToken":"asdk;lf;dsjlkfd","SteamID":1234}}

src/fixtures/maFiles/compat/missing-account-name/manifest.json

@@ -0,0 +1 @@
+{"encrypted":false,"first_run":true,"entries":[{"encryption_iv":null,"encryption_salt":null,"filename":"1234.maFile","steamid":1234}],"periodic_checking":false,"periodic_checking_interval":5,"periodic_checking_checkall":false,"auto_confirm_market_transactions":false,"auto_confirm_trades":false}

src/main.rs

@@ -1,5 +1,5 @@
 extern crate rpassword;
-use clap::{crate_version, App, Arg, Shell};
+use clap::{crate_version, App, Arg, ArgMatches, Shell};
 use log::*;
 use std::str::FromStr;
 use std::{
@@ -12,6 +12,8 @@ use steamguard::{
 	SteamGuardAccount, UserLogin,
 };
 
+use crate::accountmanager::ManifestAccountLoadError;
+
 #[macro_use]
 extern crate lazy_static;
 #[macro_use]
@@ -199,8 +201,14 @@ fn run() -> anyhow::Result<()> {
 	manifest.submit_passkey(passkey);
 
 	loop {
-		match manifest.load_accounts() {
-			Ok(_) => break,
+		match manifest.auto_upgrade() {
+			Ok(upgraded) => {
+				if upgraded {
+					info!("Manifest auto-upgraded");
+					manifest.save()?;
+				}
+				break;
+			}
 			Err(
 				accountmanager::ManifestAccountLoadError::MissingPasskey
 				| accountmanager::ManifestAccountLoadError::IncorrectPasskey,
@@ -352,12 +360,14 @@ fn run() -> anyhow::Result<()> {
 			}
 			manifest.submit_passkey(passkey);
 		}
+		manifest.load_accounts()?;
 		for entry in &mut manifest.entries {
 			entry.encryption = Some(accountmanager::EntryEncryptionParams::generate());
 		}
 		manifest.save()?;
 		return Ok(());
 	} else if matches.is_present("decrypt") {
+		manifest.load_accounts()?;
 		for entry in &mut manifest.entries {
 			entry.encryption = None;
 		}
@@ -366,27 +376,7 @@ fn run() -> anyhow::Result<()> {
 		return Ok(());
 	}
 
-	let mut selected_accounts: Vec<Arc<Mutex<SteamGuardAccount>>> = 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 entry in &manifest.entries {
-			selected_accounts.push(manifest.get_account(&entry.account_name).unwrap().clone());
-		}
-	} else {
-		for entry in &manifest.entries {
-			if !matches.is_present("username") {
-				selected_accounts.push(manifest.get_account(&entry.account_name).unwrap().clone());
-				break;
-			}
-			if matches.value_of("username").unwrap() == entry.account_name {
-				selected_accounts.push(manifest.get_account(&entry.account_name).unwrap().clone());
-				break;
-			}
-		}
-	}
+	let mut selected_accounts = get_selected_accounts(&matches, &mut manifest)?;
 
 	debug!(
 		"selected accounts: {:?}",
@@ -547,6 +537,34 @@ fn run() -> anyhow::Result<()> {
 	Ok(())
 }
 
+fn get_selected_accounts(
+	matches: &ArgMatches,
+	manifest: &mut accountmanager::Manifest,
+) -> anyhow::Result<Vec<Arc<Mutex<SteamGuardAccount>>>, ManifestAccountLoadError> {
+	let mut selected_accounts: Vec<Arc<Mutex<SteamGuardAccount>>> = vec![];
+
+	if matches.is_present("all") {
+		manifest.load_accounts()?;
+		for entry in &manifest.entries {
+			selected_accounts.push(manifest.get_account(&entry.account_name).unwrap().clone());
+		}
+	} else {
+		let entry = if matches.is_present("username") {
+			manifest.get_entry(&matches.value_of("username").unwrap().into())
+		} else {
+			manifest
+				.entries
+				.first()
+				.ok_or(ManifestAccountLoadError::MissingManifestEntry)
+		}?;
+
+		let account_name = entry.account_name.clone();
+		let account = manifest.get_or_load_account(&account_name)?;
+		selected_accounts.push(account);
+	}
+	return Ok(selected_accounts);
+}
+
 fn do_login(account: &mut SteamGuardAccount) -> anyhow::Result<()> {
 	if account.account_name.len() > 0 {
 		println!("Username: {}", account.account_name);