add authenticator 2fa support

6 files changed, 249 insertions, 55 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 021f4ce..c581bbf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Added
+
+* Authenticator-based two-step login is now supported.
+
 ### Fixed
 
 * Correctly handle password retries when entering an invalid password on the
diff --git a/src/actions.rs b/src/actions.rs
index 172dba6..c63191f 100644
--- a/src/actions.rs
+++ b/src/actions.rs
@@ -3,6 +3,8 @@ use crate::prelude::*;
 pub async fn login(
     email: &str,
     password: &crate::locked::Password,
+    two_factor_token: Option<&str>,
+    two_factor_provider: Option<crate::api::TwoFactorProviderType>,
 ) -> Result<(String, String, u32, String, crate::locked::Keys)> {
     let config = crate::config::Config::load_async().await?;
     let client =
@@ -13,7 +15,12 @@ pub async fn login(
         crate::identity::Identity::new(email, password, iterations)?;
 
     let (access_token, refresh_token, protected_key) = client
-        .login(&identity.email, &identity.master_password_hash)
+        .login(
+            &identity.email,
+            &identity.master_password_hash,
+            two_factor_token,
+            two_factor_provider,
+        )
         .await?;
     let master_keys = crate::cipherstring::CipherString::new(&protected_key)?
         .decrypt_locked_symmetric(&identity.keys)?;
diff --git a/src/api.rs b/src/api.rs
index c434069..813043a 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -4,6 +4,34 @@ use crate::json::{
     DeserializeJsonWithPath as _, DeserializeJsonWithPathAsync as _,
 };
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum TwoFactorProviderType {
+    Authenticator = 0,
+    Email = 1,
+    Duo = 2,
+    Yubikey = 3,
+    U2f = 4,
+    Remember = 5,
+    OrganizationDuo = 6,
+}
+
+impl std::convert::TryFrom<u32> for TwoFactorProviderType {
+    type Error = Error;
+
+    fn try_from(ty: u32) -> Result<Self> {
+        match ty {
+            0 => Ok(Self::Authenticator),
+            1 => Ok(Self::Email),
+            2 => Ok(Self::Duo),
+            3 => Ok(Self::Yubikey),
+            4 => Ok(Self::U2f),
+            5 => Ok(Self::Remember),
+            6 => Ok(Self::OrganizationDuo),
+            _ => Err(Error::InvalidTwoFactorProvider { ty }),
+        }
+    }
+}
+
 #[derive(serde::Serialize, Debug)]
 struct PreloginReq {
     email: String,
@@ -32,6 +60,10 @@ struct ConnectPasswordReq {
     device_name: String,
     #[serde(rename = "devicePushToken")]
     device_push_token: String,
+    #[serde(rename = "twoFactorToken")]
+    two_factor_token: Option<String>,
+    #[serde(rename = "twoFactorProvider")]
+    two_factor_provider: Option<u32>,
 }
 
 #[derive(serde::Deserialize, Debug)]
@@ -405,6 +437,8 @@ impl Client {
         &self,
         email: &str,
         master_password_hash: &crate::locked::PasswordHash,
+        two_factor_token: Option<&str>,
+        two_factor_provider: Option<TwoFactorProviderType>,
     ) -> Result<(String, String, String)> {
         let connect_req = ConnectPasswordReq {
             grant_type: "password".to_string(),
@@ -418,6 +452,9 @@ impl Client {
                 .to_string(),
             device_name: "test cli".to_string(),
             device_push_token: "".to_string(),
+            two_factor_token: two_factor_token
+                .map(std::string::ToString::to_string),
+            two_factor_provider: two_factor_provider.map(|ty| ty as u32),
         };
         let client = reqwest::Client::new();
         let res = client
@@ -869,8 +906,16 @@ fn classify_login_error(error_res: &ConnectErrorRes, code: u16) -> Error {
             "Two factor required." => {
                 match error_res.two_factor_providers.as_ref() {
                     Some(providers) => {
-                        return Error::TwoFactorRequired {
-                            providers: providers.clone(),
+                        let providers: Result<_> = providers
+                            .iter()
+                            .copied()
+                            .map(std::convert::TryInto::try_into)
+                            .collect();
+                        return match providers {
+                            Ok(providers) => {
+                                Error::TwoFactorRequired { providers }
+                            }
+                            Err(e) => e,
                         };
                     }
                     _ => {}
@@ -884,10 +929,17 @@ fn classify_login_error(error_res: &ConnectErrorRes, code: u16) -> Error {
             if error_res.error_description == "" {
                 if let Some(error_model) = error_res.error_model.as_ref() {
                     match error_model.message.as_str() {
-                        "Username or password is incorrect. Try again" => {
+                        "Username or password is incorrect. Try again"
+                        | "TOTP code is not a number" => {
                             return Error::IncorrectPassword;
                         }
-                        _ => {}
+                        s => {
+                            if s.starts_with(
+                                "Invalid TOTP code! Server time: ",
+                            ) {
+                                return Error::IncorrectPassword;
+                            }
+                        }
                     }
                 }
             }
diff --git a/src/bin/rbw-agent/actions.rs b/src/bin/rbw-agent/actions.rs
index 0b05829..d021bda 100644
--- a/src/bin/rbw-agent/actions.rs
+++ b/src/bin/rbw-agent/actions.rs
@@ -5,7 +5,7 @@ pub async fn login(
     state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
     tty: Option<&str>,
 ) -> anyhow::Result<()> {
-    let mut db = load_db().await.unwrap_or_else(|_| rbw::db::Db::new());
+    let db = load_db().await.unwrap_or_else(|_| rbw::db::Db::new());
 
     if db.needs_login() {
         let url_str = config_base_url().await?;
@@ -36,8 +36,7 @@ pub async fn login(
             )
             .await
             .context("failed to read password from pinentry")?;
-            let res = rbw::actions::login(&email, &password).await;
-            match res {
+            match rbw::actions::login(&email, &password, None, None).await {
                 Ok((
                     access_token,
                     refresh_token,
@@ -45,50 +44,53 @@ pub async fn login(
                     protected_key,
                     _,
                 )) => {
-                    db.access_token = Some(access_token);
-                    db.refresh_token = Some(refresh_token);
-                    db.iterations = Some(iterations);
-                    db.protected_key = Some(protected_key.clone());
-                    save_db(&db).await?;
-
-                    sync(sock, false).await?;
-                    db = load_db().await?;
-
-                    let protected_private_key =
-                        if let Some(protected_private_key) =
-                            db.protected_private_key
-                        {
-                            protected_private_key
-                        } else {
-                            return Err(anyhow::anyhow!(
-                                "failed to find protected private key in db"
-                            ));
-                        };
-
-                    let res = rbw::actions::unlock(
-                        &email,
-                        &password,
+                    login_success(
+                        sock,
+                        state,
+                        access_token,
+                        refresh_token,
                         iterations,
-                        &protected_key,
-                        &protected_private_key,
-                        &db.protected_org_keys,
+                        protected_key,
+                        password,
+                        db,
+                        email,
                     )
-                    .await;
-
-                    match res {
-                        Ok((keys, org_keys)) => {
-                            let mut state = state.write().await;
-                            state.priv_key = Some(keys);
-                            state.org_keys = Some(org_keys);
-                        }
-                        Err(e) => {
-                            return Err(e)
-                                .context("failed to unlock database")
-                        }
-                    }
-
+                    .await?;
                     break;
                 }
+                Err(rbw::error::Error::TwoFactorRequired { providers }) => {
+                    if providers.contains(
+                        &rbw::api::TwoFactorProviderType::Authenticator,
+                    ) {
+                        let (
+                            access_token,
+                            refresh_token,
+                            iterations,
+                            protected_key,
+                        ) = two_factor(
+                            tty,
+                            &email,
+                            &password,
+                            rbw::api::TwoFactorProviderType::Authenticator,
+                        )
+                        .await?;
+                        login_success(
+                            sock,
+                            state,
+                            access_token,
+                            refresh_token,
+                            iterations,
+                            protected_key,
+                            password,
+                            db,
+                            email,
+                        )
+                        .await?;
+                        break;
+                    } else {
+                        return Err(anyhow::anyhow!("TODO"));
+                    }
+                }
                 Err(rbw::error::Error::IncorrectPassword) => {
                     if i == 3 {
                         return Err(rbw::error::Error::IncorrectPassword)
@@ -112,6 +114,119 @@ pub async fn login(
     Ok(())
 }
 
+async fn two_factor(
+    tty: Option<&str>,
+    email: &str,
+    password: &rbw::locked::Password,
+    provider: rbw::api::TwoFactorProviderType,
+) -> anyhow::Result<(String, String, u32, String)> {
+    for i in 1_u8..=3 {
+        let err = if i > 1 {
+            Some(format!("Incorrect code (attempt {}/3)", i))
+        } else {
+            None
+        };
+        let code = rbw::pinentry::getpin(
+            "Authenticator App",
+            "Enter the 6 digit verification code from your authenticator app.",
+            err.as_deref(),
+            tty,
+        )
+        .await
+        .context("failed to read code from pinentry")?;
+        let code = std::str::from_utf8(code.password())
+            .context("code was not valid utf8")?;
+        match rbw::actions::login(
+            &email,
+            &password,
+            Some(code),
+            Some(provider),
+        )
+        .await
+        {
+            Ok((
+                access_token,
+                refresh_token,
+                iterations,
+                protected_key,
+                _,
+            )) => {
+                return Ok((
+                    access_token,
+                    refresh_token,
+                    iterations,
+                    protected_key,
+                ))
+            }
+            Err(rbw::error::Error::IncorrectPassword) => {
+                if i == 3 {
+                    return Err(rbw::error::Error::IncorrectPassword)
+                        .context("failed to log in to bitwarden instance");
+                } else {
+                    continue;
+                }
+            }
+            Err(e) => {
+                return Err(e)
+                    .context("failed to log in to bitwarden instance")
+            }
+        }
+    }
+
+    unreachable!()
+}
+
+async fn login_success(
+    sock: &mut crate::sock::Sock,
+    state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
+    access_token: String,
+    refresh_token: String,
+    iterations: u32,
+    protected_key: String,
+    password: rbw::locked::Password,
+    mut db: rbw::db::Db,
+    email: String,
+) -> anyhow::Result<()> {
+    db.access_token = Some(access_token.to_string());
+    db.refresh_token = Some(refresh_token.to_string());
+    db.iterations = Some(iterations);
+    db.protected_key = Some(protected_key.to_string());
+    save_db(&db).await?;
+
+    sync(sock, false).await?;
+    let db = load_db().await?;
+
+    let protected_private_key =
+        if let Some(protected_private_key) = db.protected_private_key {
+            protected_private_key
+        } else {
+            return Err(anyhow::anyhow!(
+                "failed to find protected private key in db"
+            ));
+        };
+
+    let res = rbw::actions::unlock(
+        &email,
+        &password,
+        iterations,
+        &protected_key,
+        &protected_private_key,
+        &db.protected_org_keys,
+    )
+    .await;
+
+    match res {
+        Ok((keys, org_keys)) => {
+            let mut state = state.write().await;
+            state.priv_key = Some(keys);
+            state.org_keys = Some(org_keys);
+        }
+        Err(e) => return Err(e).context("failed to unlock database"),
+    }
+
+    Ok(())
+}
+
 pub async fn unlock(
     sock: &mut crate::sock::Sock,
     state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
@@ -159,7 +274,7 @@ pub async fn unlock(
             )
             .await
             .context("failed to read password from pinentry")?;
-            let res = rbw::actions::unlock(
+            match rbw::actions::unlock(
                 &email,
                 &password,
                 iterations,
@@ -167,12 +282,10 @@ pub async fn unlock(
                 &protected_private_key,
                 &db.protected_org_keys,
             )
-            .await;
-            match res {
+            .await
+            {
                 Ok((keys, org_keys)) => {
-                    let mut state = state.write().await;
-                    state.priv_key = Some(keys);
-                    state.org_keys = Some(org_keys);
+                    unlock_success(state, keys, org_keys).await?;
                     break;
                 }
                 Err(rbw::error::Error::IncorrectPassword) => {
@@ -193,6 +306,17 @@ pub async fn unlock(
     Ok(())
 }
 
+async fn unlock_success(
+    state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
+    keys: rbw::locked::Keys,
+    org_keys: std::collections::HashMap<String, rbw::locked::Keys>,
+) -> anyhow::Result<()> {
+    let mut state = state.write().await;
+    state.priv_key = Some(keys);
+    state.org_keys = Some(org_keys);
+    Ok(())
+}
+
 pub async fn lock(
     sock: &mut crate::sock::Sock,
     state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
diff --git a/src/bin/rbw-agent/main.rs b/src/bin/rbw-agent/main.rs
index 5e476f8..1a7c382 100644
--- a/src/bin/rbw-agent/main.rs
+++ b/src/bin/rbw-agent/main.rs
@@ -1,3 +1,5 @@
+#![allow(clippy::too_many_arguments)]
+
 use anyhow::Context as _;
 
 mod actions;
diff --git a/src/error.rs b/src/error.rs
index ea57665..5a6eef4 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -49,6 +49,9 @@ pub enum Error {
     #[snafu(display("invalid mac"))]
     InvalidMac,
 
+    #[snafu(display("invalid two factor provider type: {}", ty))]
+    InvalidTwoFactorProvider { ty: u32 },
+
     #[snafu(display("failed to parse JSON"))]
     JSON {
         source: serde_path_to_error::Error<serde_json::Error>,
@@ -157,7 +160,9 @@ pub enum Error {
     Spawn { source: tokio::io::Error },
 
     #[snafu(display("two factor required"))]
-    TwoFactorRequired { providers: Vec<u32> },
+    TwoFactorRequired {
+        providers: Vec<crate::api::TwoFactorProviderType>,
+    },
 
     #[snafu(display("error writing to pinentry stdin"))]
     WriteStdin { source: tokio::io::Error },