expand the protocol to allow decrypting with organization keys

not implemented yet, just changing the interface

5 files changed, 75 insertions, 33 deletions

diff --git a/src/bin/rbw-agent/actions.rs b/src/bin/rbw-agent/actions.rs
index 52c34fa..afe76c8 100644
--- a/src/bin/rbw-agent/actions.rs
+++ b/src/bin/rbw-agent/actions.rs
@@ -211,9 +211,10 @@ pub async fn decrypt(
     sock: &mut crate::sock::Sock,
     state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
     cipherstring: &str,
+    org_id: Option<&str>,
 ) -> anyhow::Result<()> {
     let state = state.read().await;
-    let keys = if let Some(keys) = &state.priv_key {
+    let keys = if let Some(keys) = state.key(org_id) {
         keys
     } else {
         return Err(anyhow::anyhow!(
@@ -238,9 +239,10 @@ pub async fn encrypt(
     sock: &mut crate::sock::Sock,
     state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
     plaintext: &str,
+    org_id: Option<&str>,
 ) -> anyhow::Result<()> {
     let state = state.read().await;
-    let keys = if let Some(keys) = &state.priv_key {
+    let keys = if let Some(keys) = state.key(org_id) {
         keys
     } else {
         return Err(anyhow::anyhow!(
diff --git a/src/bin/rbw-agent/agent.rs b/src/bin/rbw-agent/agent.rs
index 72d6077..0094f7d 100644
--- a/src/bin/rbw-agent/agent.rs
+++ b/src/bin/rbw-agent/agent.rs
@@ -9,10 +9,18 @@ pub enum TimeoutEvent {
 
 pub struct State {
     pub priv_key: Option<rbw::locked::Keys>,
+    pub org_keys: std::collections::HashMap<String, rbw::locked::Keys>,
     pub timeout_chan: tokio::sync::mpsc::UnboundedSender<TimeoutEvent>,
 }
 
 impl State {
+    pub fn key(&self, org_id: Option<&str>) -> Option<&rbw::locked::Keys> {
+        match org_id {
+            Some(id) => self.org_keys.get(id),
+            None => self.priv_key.as_ref(),
+        }
+    }
+
     pub fn needs_unlock(&self) -> bool {
         self.priv_key.is_none()
     }
@@ -24,6 +32,7 @@ impl State {
 
     pub fn clear(&mut self) {
         self.priv_key = None;
+        self.org_keys = Default::default();
         // no real better option to unwrap here
         self.timeout_chan.send(TimeoutEvent::Clear).unwrap();
     }
@@ -49,6 +58,7 @@ impl Agent {
             timeout_chan: r,
             state: std::sync::Arc::new(tokio::sync::RwLock::new(State {
                 priv_key: None,
+                org_keys: Default::default(),
                 timeout_chan: w,
             })),
         })
@@ -146,13 +156,27 @@ async fn handle_request(
             crate::actions::sync(sock).await?;
             false
         }
-        rbw::protocol::Action::Decrypt { cipherstring } => {
-            crate::actions::decrypt(sock, state.clone(), &cipherstring)
-                .await?;
+        rbw::protocol::Action::Decrypt {
+            cipherstring,
+            org_id,
+        } => {
+            crate::actions::decrypt(
+                sock,
+                state.clone(),
+                &cipherstring,
+                org_id.as_deref(),
+            )
+            .await?;
             true
         }
-        rbw::protocol::Action::Encrypt { plaintext } => {
-            crate::actions::encrypt(sock, state.clone(), &plaintext).await?;
+        rbw::protocol::Action::Encrypt { plaintext, org_id } => {
+            crate::actions::encrypt(
+                sock,
+                state.clone(),
+                &plaintext,
+                org_id.as_deref(),
+            )
+            .await?;
             true
         }
         rbw::protocol::Action::Quit => std::process::exit(0),
diff --git a/src/bin/rbw/actions.rs b/src/bin/rbw/actions.rs
index 399b7f3..6de2fc6 100644
--- a/src/bin/rbw/actions.rs
+++ b/src/bin/rbw/actions.rs
@@ -42,13 +42,17 @@ pub fn quit() -> anyhow::Result<()> {
     }
 }
 
-pub fn decrypt(cipherstring: &str) -> anyhow::Result<String> {
+pub fn decrypt(
+    cipherstring: &str,
+    org_id: Option<&str>,
+) -> anyhow::Result<String> {
     let mut sock = crate::sock::Sock::connect()
         .context("failed to connect to rbw-agent")?;
     sock.send(&rbw::protocol::Request {
         tty: std::env::var("TTY").ok(),
         action: rbw::protocol::Action::Decrypt {
             cipherstring: cipherstring.to_string(),
+            org_id: org_id.map(std::string::ToString::to_string),
         },
     })?;
 
@@ -62,13 +66,17 @@ pub fn decrypt(cipherstring: &str) -> anyhow::Result<String> {
     }
 }
 
-pub fn encrypt(plaintext: &str) -> anyhow::Result<String> {
+pub fn encrypt(
+    plaintext: &str,
+    org_id: Option<&str>,
+) -> anyhow::Result<String> {
     let mut sock = crate::sock::Sock::connect()
         .context("failed to connect to rbw-agent")?;
     sock.send(&rbw::protocol::Request {
         tty: std::env::var("TTY").ok(),
         action: rbw::protocol::Action::Encrypt {
             plaintext: plaintext.to_string(),
+            org_id: org_id.map(std::string::ToString::to_string),
         },
     })?;
 
diff --git a/src/bin/rbw/commands.rs b/src/bin/rbw/commands.rs
index e05f2eb..88b1fae 100644
--- a/src/bin/rbw/commands.rs
+++ b/src/bin/rbw/commands.rs
@@ -189,24 +189,24 @@ pub fn add(
     let mut access_token = db.access_token.as_ref().unwrap().clone();
     let refresh_token = db.refresh_token.as_ref().unwrap();
 
-    let name = crate::actions::encrypt(name)?;
+    let name = crate::actions::encrypt(name, None)?;
 
     let username = username
-        .map(|username| crate::actions::encrypt(username))
+        .map(|username| crate::actions::encrypt(username, None))
         .transpose()?;
 
     let contents = rbw::edit::edit("", HELP)?;
 
     let (password, notes) = parse_editor(&contents);
     let password = password
-        .map(|password| crate::actions::encrypt(&password))
+        .map(|password| crate::actions::encrypt(&password, None))
         .transpose()?;
     let notes = notes
-        .map(|notes| crate::actions::encrypt(&notes))
+        .map(|notes| crate::actions::encrypt(&notes, None))
         .transpose()?;
     let uris: Vec<String> = uris
         .iter()
-        .map(|uri| crate::actions::encrypt(&uri))
+        .map(|uri| crate::actions::encrypt(&uri, None))
         .collect::<anyhow::Result<_>>()?;
 
     let mut folder_id = None;
@@ -222,7 +222,7 @@ pub fn add(
         let folders: Vec<(String, String)> = folders
             .iter()
             .cloned()
-            .map(|(id, name)| Ok((id, crate::actions::decrypt(&name)?)))
+            .map(|(id, name)| Ok((id, crate::actions::decrypt(&name, None)?)))
             .collect::<anyhow::Result<_>>()?;
 
         for (id, name) in folders {
@@ -234,7 +234,7 @@ pub fn add(
             let (new_access_token, id) = rbw::actions::create_folder(
                 &access_token,
                 &refresh_token,
-                &crate::actions::encrypt(folder_name)?,
+                &crate::actions::encrypt(folder_name, None)?,
             )?;
             if let Some(new_access_token) = new_access_token {
                 access_token = new_access_token.clone();
@@ -285,14 +285,14 @@ pub fn generate(
         let mut access_token = db.access_token.as_ref().unwrap().clone();
         let refresh_token = db.refresh_token.as_ref().unwrap();
 
-        let name = crate::actions::encrypt(name)?;
+        let name = crate::actions::encrypt(name, None)?;
         let username = username
-            .map(|username| crate::actions::encrypt(username))
+            .map(|username| crate::actions::encrypt(username, None))
             .transpose()?;
-        let password = crate::actions::encrypt(&password)?;
+        let password = crate::actions::encrypt(&password, None)?;
         let uris: Vec<String> = uris
             .iter()
-            .map(|uri| crate::actions::encrypt(&uri))
+            .map(|uri| crate::actions::encrypt(&uri, None))
             .collect::<anyhow::Result<_>>()?;
 
         let mut folder_id = None;
@@ -308,7 +308,9 @@ pub fn generate(
             let folders: Vec<(String, String)> = folders
                 .iter()
                 .cloned()
-                .map(|(id, name)| Ok((id, crate::actions::decrypt(&name)?)))
+                .map(|(id, name)| {
+                    Ok((id, crate::actions::decrypt(&name, None)?))
+                })
                 .collect::<anyhow::Result<_>>()?;
 
             for (id, name) in folders {
@@ -320,7 +322,7 @@ pub fn generate(
                 let (new_access_token, id) = rbw::actions::create_folder(
                     &access_token,
                     &refresh_token,
-                    &crate::actions::encrypt(folder_name)?,
+                    &crate::actions::encrypt(folder_name, None)?,
                 )?;
                 if let Some(new_access_token) = new_access_token {
                     access_token = new_access_token.clone();
@@ -381,10 +383,10 @@ pub fn edit(name: &str, username: Option<&str>) -> anyhow::Result<()> {
 
     let (password, notes) = parse_editor(&contents);
     let password = password
-        .map(|password| crate::actions::encrypt(&password))
+        .map(|password| crate::actions::encrypt(&password, None))
         .transpose()?;
     let notes = notes
-        .map(|notes| crate::actions::encrypt(&notes))
+        .map(|notes| crate::actions::encrypt(&notes, None))
         .transpose()?;
     let mut history = entry.history.clone();
     let new_history_entry = rbw::db::HistoryEntry {
@@ -660,7 +662,7 @@ fn decrypt_cipher(entry: &rbw::db::Entry) -> anyhow::Result<DecryptedCipher> {
     let folder = entry
         .folder
         .as_ref()
-        .map(|folder| crate::actions::decrypt(folder))
+        .map(|folder| crate::actions::decrypt(folder, None))
         .transpose();
     let folder = match folder {
         Ok(folder) => folder,
@@ -672,7 +674,7 @@ fn decrypt_cipher(entry: &rbw::db::Entry) -> anyhow::Result<DecryptedCipher> {
     let username = entry
         .username
         .as_ref()
-        .map(|username| crate::actions::decrypt(username))
+        .map(|username| crate::actions::decrypt(username, None))
         .transpose();
     let username = match username {
         Ok(username) => username,
@@ -684,7 +686,7 @@ fn decrypt_cipher(entry: &rbw::db::Entry) -> anyhow::Result<DecryptedCipher> {
     let password = entry
         .password
         .as_ref()
-        .map(|password| crate::actions::decrypt(password))
+        .map(|password| crate::actions::decrypt(password, None))
         .transpose();
     let password = match password {
         Ok(password) => password,
@@ -696,7 +698,7 @@ fn decrypt_cipher(entry: &rbw::db::Entry) -> anyhow::Result<DecryptedCipher> {
     let notes = entry
         .notes
         .as_ref()
-        .map(|notes| crate::actions::decrypt(notes))
+        .map(|notes| crate::actions::decrypt(notes, None))
         .transpose();
     let notes = match notes {
         Ok(notes) => notes,
@@ -711,14 +713,14 @@ fn decrypt_cipher(entry: &rbw::db::Entry) -> anyhow::Result<DecryptedCipher> {
         .map(|entry| {
             Ok(DecryptedHistoryEntry {
                 last_used_date: entry.last_used_date.clone(),
-                password: crate::actions::decrypt(&entry.password)?,
+                password: crate::actions::decrypt(&entry.password, None)?,
             })
         })
         .collect::<anyhow::Result<_>>()?;
     Ok(DecryptedCipher {
         id: entry.id.clone(),
         folder,
-        name: crate::actions::decrypt(&entry.name)?,
+        name: crate::actions::decrypt(&entry.name, None)?,
         username,
         password,
         notes,
diff --git a/src/protocol.rs b/src/protocol.rs
index 2ef99cd..0ab594c 100644
--- a/src/protocol.rs
+++ b/src/protocol.rs
@@ -1,4 +1,4 @@
-pub const VERSION: u32 = 1;
+pub const VERSION: u32 = 2;
 
 #[derive(serde::Serialize, serde::Deserialize, Debug)]
 pub struct Request {
@@ -13,8 +13,14 @@ pub enum Action {
     Unlock,
     Lock,
     Sync,
-    Decrypt { cipherstring: String },
-    Encrypt { plaintext: String },
+    Decrypt {
+        cipherstring: String,
+        org_id: Option<String>,
+    },
+    Encrypt {
+        plaintext: String,
+        org_id: Option<String>,
+    },
     Quit,
     Version,
 }