SM-866: Implement bws State

crates/bws/Cargo.toml

@@ -39,6 +39,7 @@ thiserror = "1.0.50"
 tokio = { version = "1.34.0", features = ["rt-multi-thread", "macros"] }
 toml = "0.8.8"
 uuid = { version = "^1.6.1", features = ["serde"] }
+sha2 = ">=0.10.6, <0.11"
 
 bitwarden = { path = "../bitwarden", version = "0.3.1", features = ["secrets"] }
 

crates/bws/src/main.rs

@@ -13,7 +13,6 @@ use bitwarden::{
             SecretIdentifiersRequest, SecretPutRequest, SecretsDeleteRequest, SecretsGetRequest,
         },
     },
-    state::StateManager,
 };
 use clap::{ArgGroup, CommandFactory, Parser, Subcommand};
 use clap_complete::Shell;
@@ -26,7 +25,7 @@ mod state;
 
 use config::ProfileKey;
 use render::{serialize_response, Color, Output};
-use serde_json::json;
+use state::State;
 use uuid::Uuid;
 
 #[derive(Parser, Debug)]
@@ -333,24 +332,36 @@ async fn process_commands() -> Result<()> {
         true,
     );
 
-    let mut state = StateManager::new(&state_file_path)?;
-    let client_state = state.get_client_state();
-    let valid_token = client_state.token_is_valid();
+    let mut state = State::new(&state_file_path, access_token.clone())?;
+    let mut valid_token = false;
+    let client_state = match state.get() {
+        Some(state) => match state {
+            Ok(client_state) => {
+                valid_token = client_state.token_is_valid();
+                Some(client_state)
+            }
+            Err(_) => {
+                println!("Error decrypting the client state. Proceeding without state, it might be overwritten!");
+                None
+            }
+        },
+        None => None,
+    };
 
-    let mut client = bitwarden::Client::new(settings, Some(client_state));
+    let mut client = bitwarden::Client::new(settings, client_state);
 
-    // TODO: Remove println! commands below
     if !valid_token {
-        println!("calling access_token_login...");
         let _ = client
             .auth()
             .login_access_token(&AccessTokenLoginRequest { access_token })
             .await?;
 
-        state.data = json!(client.get_client_state());
-        println!("state data: {:?}", state.data);
-        let r = state.save(&state_file_path);
-        println!("save result: {:?}", r);
+        if state.upsert(client.get_client_state()).is_err() {
+            println!("Failure to update the in-memory state.")
+        }
+        if state.save(&state_file_path).is_err() {
+            println!("Failure to save the state.")
+        }
     }
 
     let organization_id = match client.get_access_token_organization() {

crates/bws/src/state.rs

@@ -1,15 +1,86 @@
-use std::path::PathBuf;
+use std::{
+    collections::HashMap,
+    path::{Path, PathBuf},
+};
 
+use bitwarden::{
+    client::{AccessToken, ClientState},
+    crypto::{EncString, KeyDecryptable, KeyEncryptable},
+    error::{Error, Result},
+    state::StateManager,
+};
 use directories::BaseDirs;
-use serde::{Deserialize, Serialize};
+use serde_json::json;
+use sha2::Digest;
+
+type AccessTokenHash = String;
+type EncClientState = EncString;
 
 pub(crate) const ROOT_DIRECTORY: &str = ".bws";
 pub(crate) const STATE_DIRECTORY: &str = "state";
 pub(crate) const FILENAME: &str = "state";
 
-#[derive(Serialize, Deserialize)]
-struct State {
-    version: u32,
+pub struct State {
+    state_manager: StateManager,
+    data: HashMap<AccessTokenHash, EncClientState>,
+    access_token: AccessToken,
+    access_token_hash: AccessTokenHash,
+}
+
+impl State {
+    pub fn new(path: &Path, access_token: String) -> Result<Self> {
+        let state_manager = StateManager::new(path)?;
+        let data: HashMap<AccessTokenHash, EncClientState> = if state_manager.has_data() {
+            serde_json::from_str(state_manager.data.to_string().as_str())?
+        } else {
+            Default::default()
+        };
+        let access_token_hash: String = format!(
+            "{:X}",
+            sha2::Sha256::new()
+                .chain_update(access_token.clone())
+                .finalize()
+        );
+        let access_token: AccessToken = access_token.parse()?;
+
+        Ok(Self {
+            state_manager,
+            data,
+            access_token,
+            access_token_hash,
+        })
+    }
+
+    pub fn get(&self) -> Option<Result<ClientState>> {
+        match self.data.get(&self.access_token_hash) {
+            Some(encrypted_data) => {
+                let decrypted_data: Result<String> =
+                    encrypted_data.decrypt_with_key(&self.access_token.encryption_key);
+                match decrypted_data {
+                    Ok(decrypted_data) => match serde_json::from_str(decrypted_data.as_str()) {
+                        Ok(state) => Some(Ok(state)),
+                        Err(e) => Some(Err(Error::Serde(e))),
+                    },
+                    Err(e) => Some(Err(e)),
+                }
+            }
+            None => None,
+        }
+    }
+
+    pub fn upsert(&mut self, new_state: ClientState) -> Result<()> {
+        let serialized_state = json!(new_state).to_string();
+        let enc_state = serialized_state.encrypt_with_key(&self.access_token.encryption_key)?;
+        self.data.insert(self.access_token_hash.clone(), enc_state);
+        self.state_manager.data = json!(self.data);
+
+        Ok(())
+    }
+
+    pub fn save(&mut self, path: &Path) -> Result<()> {
+        self.state_manager.data = json!(self.data);
+        self.state_manager.save(path)
+    }
 }
 
 pub(crate) fn get_state_file_path(