From fdeafff94198dce1f48378f385c8ca591b63cff3 Mon Sep 17 00:00:00 2001
From: Carson McManus <carson.mcmanus1@gmail.com>
Date: Sun, 21 Jul 2024 10:18:20 -0400
Subject: [PATCH] add `status` subcommand

---
 src/commands.rs                      |  3 +
 src/commands/status.rs               | 85 ++++++++++++++++++++++++++++
 src/main.rs                          |  1 +
 steamguard/src/steamapi/twofactor.rs |  4 +-
 4 files changed, 91 insertions(+), 2 deletions(-)
 create mode 100644 src/commands/status.rs

diff --git a/src/commands.rs b/src/commands.rs
index 72e94cae..2d303c9e 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -3,6 +3,7 @@ use std::sync::{Arc, Mutex};
 use clap::{Parser, Subcommand, ValueEnum};
 use clap_complete::Shell;
 use secrecy::SecretString;
+use status::StatusCommand;
 use std::str::FromStr;
 use steamguard::{transport::Transport, SteamGuardAccount};
 
@@ -20,6 +21,7 @@ pub mod qr;
 pub mod qr_login;
 pub mod remove;
 pub mod setup;
+pub mod status;
 
 pub use code::CodeCommand;
 pub use completions::CompletionsCommand;
@@ -175,6 +177,7 @@ pub(crate) enum Subcommands {
 	#[cfg(feature = "qr")]
 	Qr(QrCommand),
 	QrLogin(QrLoginCommand),
+	Status(StatusCommand),
 }
 
 #[derive(Debug, Clone, Copy, ValueEnum)]
diff --git a/src/commands/status.rs b/src/commands/status.rs
new file mode 100644
index 00000000..c9d1a386
--- /dev/null
+++ b/src/commands/status.rs
@@ -0,0 +1,85 @@
+use clap::Parser;
+use log::*;
+use steamguard::{
+	protobufs::service_twofactor::CTwoFactor_Status_Request,
+	steamapi::TwoFactorClient,
+	transport::{Transport, TransportError},
+	SteamGuardAccount,
+};
+
+use super::AccountCommand;
+
+#[derive(Debug, Clone, Parser)]
+#[clap(about = "Query and print the 2FA status of an account.")]
+pub struct StatusCommand;
+
+impl<T> AccountCommand<T> for StatusCommand
+where
+	T: Transport + Clone,
+{
+	fn execute(
+		&self,
+		transport: T,
+		manager: &mut crate::AccountManager,
+		accounts: Vec<std::sync::Arc<std::sync::Mutex<steamguard::SteamGuardAccount>>>,
+		args: &super::GlobalArgs,
+	) -> anyhow::Result<()> {
+		let client = TwoFactorClient::new(transport.clone());
+
+		for account in accounts {
+			let mut account = account.lock().unwrap();
+			match print_account_status(&mut account, &transport, args, &client) {
+				Ok(_) => {}
+				Err(e) => {
+					error!(
+						"Failed to print status for account {}: {}",
+						account.account_name, e
+					);
+				}
+			}
+		}
+
+		manager.save()?;
+
+		Ok(())
+	}
+}
+
+fn print_account_status<T>(
+	account: &mut SteamGuardAccount,
+	transport: &T,
+	args: &super::GlobalArgs,
+	client: &TwoFactorClient<T>,
+) -> anyhow::Result<()>
+where
+	T: Transport + Clone,
+{
+	if account.tokens.is_none() {
+		crate::do_login(transport.clone(), account, args.password.clone())?;
+	}
+	let Some(tokens) = account.tokens.as_ref() else {
+		bail!(
+			"No tokens found for {}. Can't query status if we aren't logged in ourselves.",
+			account.account_name
+		);
+	};
+	let mut req = CTwoFactor_Status_Request::new();
+	req.set_steamid(account.steam_id);
+	let resp = match client.query_status(req.clone(), tokens.access_token()) {
+		Ok(resp) => resp,
+		Err(TransportError::Unauthorized) => {
+			info!("Access token expired, re-logging in...");
+			crate::do_login(transport.clone(), account, args.password.clone())?;
+			let tokens = account.tokens.as_ref().unwrap();
+			client.query_status(req, tokens.access_token())?
+		}
+		Err(e) => {
+			return Err(e.into());
+		}
+	};
+	let data = resp.into_response_data();
+
+	println!("Account: {}", account.account_name);
+	println!("Status: {:#?}", data);
+	Ok(())
+}
diff --git a/src/main.rs b/src/main.rs
index cdb4d84f..167751aa 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -92,6 +92,7 @@ fn run(args: commands::Args) -> anyhow::Result<()> {
 		#[cfg(feature = "qr")]
 		Subcommands::Qr(args) => CommandType::Account(Box::new(args)),
 		Subcommands::QrLogin(args) => CommandType::Account(Box::new(args)),
+		Subcommands::Status(args) => CommandType::Account(Box::new(args)),
 	};
 
 	if let CommandType::Const(cmd) = cmd {
diff --git a/steamguard/src/steamapi/twofactor.rs b/steamguard/src/steamapi/twofactor.rs
index 0eef243d..b9bbe772 100644
--- a/steamguard/src/steamapi/twofactor.rs
+++ b/steamguard/src/steamapi/twofactor.rs
@@ -112,7 +112,7 @@ where
 		&self,
 		req: CTwoFactor_Status_Request,
 		access_token: &Jwt,
-	) -> anyhow::Result<ApiResponse<CTwoFactor_Status_Response>> {
+	) -> Result<ApiResponse<CTwoFactor_Status_Response>, TransportError> {
 		let req =
 			ApiRequest::new(SERVICE_NAME, "QueryStatus", 1, req).with_access_token(access_token);
 		let resp = self
@@ -121,7 +121,7 @@ where
 		Ok(resp)
 	}
 
-	pub fn query_time(&self) -> anyhow::Result<ApiResponse<CTwoFactor_Time_Response>> {
+	pub fn query_time(&self) -> Result<ApiResponse<CTwoFactor_Time_Response>, TransportError> {
 		let req = ApiRequest::new(SERVICE_NAME, "QueryTime", 1, CTwoFactor_Time_Request::new());
 		let resp = self
 			.transport