diff --git a/src/bin/rbw/commands.rs b/src/bin/rbw/commands.rs index 3329f76..b18c037 100644 --- a/src/bin/rbw/commands.rs +++ b/src/bin/rbw/commands.rs @@ -1943,8 +1943,21 @@ fn remove_db() -> anyhow::Result<()> { ) } -fn parse_totp_secret(secret: &str) -> anyhow::Result> { - let secret_str = if let Ok(u) = url::Url::parse(secret) { +struct TotpParams { + secret: Vec, + algorithm: String, + digits: u32, + period: u64, +} +fn decode_totp_secret(secret: String) -> anyhow::Result> { + base32::decode( + base32::Alphabet::RFC4648 { padding: false }, + &secret.replace(' ', ""), + ) + .ok_or_else(|| anyhow::anyhow!("totp secret was not valid base32")) +} +fn parse_totp_secret(secret: &str) -> anyhow::Result { + if let Ok(u) = url::Url::parse(secret) { if u.scheme() != "otpauth" { return Err(anyhow::anyhow!( "totp secret url must have otpauth scheme" @@ -1957,32 +1970,77 @@ fn parse_totp_secret(secret: &str) -> anyhow::Result> { } let query: std::collections::HashMap<_, _> = u.query_pairs().collect(); - query - .get("secret") - .ok_or_else(|| { - anyhow::anyhow!("totp secret url must have secret") - })? - .to_string() + Ok(TotpParams { + secret: decode_totp_secret(query + .get("secret") + .ok_or_else(|| { + anyhow::anyhow!("totp secret url must have secret") + })? + .to_string())?, + algorithm: match query.get("algorithm") { + Some(alg) => alg.to_string(), + None => String::from("SHA1"), + }, + digits: match query.get("digits") { + Some(dig) => { + dig.parse::().map_err(|_|{ + anyhow::anyhow!("digits parameter in totp url must be a valid integer.") + })? + } + None => 6, + }, + period: match query.get("period") { + Some(dig) => { + dig.parse::().map_err(|_|{ + anyhow::anyhow!("period parameter in totp url must be a valid integer.") + })? + } + None => 30, + } + }) } else { - secret.to_string() - }; - base32::decode( - base32::Alphabet::RFC4648 { padding: false }, - &secret_str.replace(' ', ""), - ) - .ok_or_else(|| anyhow::anyhow!("totp secret was not valid base32")) + Ok(TotpParams { + secret: decode_totp_secret(secret.to_string())?, + algorithm: String::from("SHA1"), + digits: 6, + period: 30, + }) + } } fn generate_totp(secret: &str) -> anyhow::Result { - let key = parse_totp_secret(secret)?; - Ok(totp_lite::totp_custom::( - totp_lite::DEFAULT_STEP, - 6, - &key, - std::time::SystemTime::now() - .duration_since(std::time::SystemTime::UNIX_EPOCH)? - .as_secs(), - )) + let totp_params = parse_totp_secret(secret)?; + let alg = totp_params.algorithm.as_str(); + match alg { + "SHA1" => Ok(totp_lite::totp_custom::( + totp_params.period, + totp_params.digits, + &totp_params.secret, + std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH)? + .as_secs(), + )), + "SHA256" => Ok(totp_lite::totp_custom::( + totp_params.period, + totp_params.digits, + &totp_params.secret, + std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH)? + .as_secs(), + )), + "SHA512" => Ok(totp_lite::totp_custom::( + totp_params.period, + totp_params.digits, + &totp_params.secret, + std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH)? + .as_secs(), + )), + _ => Err(anyhow::anyhow!(format!( + "{} is not a valid totp algorithm", + alg + ))), + } } fn display_field(name: &str, field: Option<&str>, clipboard: bool) -> bool {