Add SecretsStore

src/lib.rs

@@ -12,6 +12,7 @@ pub mod network;
 pub mod pcd;
 pub mod prd;
 pub mod process;
+pub mod secrets;
 pub mod signature;
 pub mod silentpayments;
 

src/secrets.rs

@@ -0,0 +1,145 @@
+use anyhow::{Result, Error};
+use crate::aes_gcm::aead::{Aead, Payload};
+use crate::aes_gcm::Nonce;
+use crate::sp_client::bitcoin::hashes::Hash;
+use crate::sp_client::silentpayments::utils::SilentPaymentAddress;
+use crate::crypto::{Aes256Gcm, AnkSharedSecretHash, KeyInit, AAD};
+use crate::log;
+use serde::ser::SerializeStruct;
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
+use std::collections::HashMap;
+use std::str::FromStr;
+
+#[derive(Debug, Clone, Default, PartialEq)]
+pub struct SecretsStore{
+    shared_secrets: HashMap<SilentPaymentAddress, AnkSharedSecretHash>,
+    unconfirmed_secrets: Vec<AnkSharedSecretHash>
+}
+
+impl Serialize for SecretsStore {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let mut temp_map = HashMap::with_capacity(self.shared_secrets.len());
+        for (key, value) in &self.shared_secrets {
+            let key_str = key.to_string(); 
+            let value_str = value.to_string();
+
+            temp_map.insert(key_str, value_str);
+        }
+
+        let mut state = serializer.serialize_struct("SecretsStore", 2)?;
+        state.serialize_field("shared_secrets", &temp_map)?;
+
+        let unconfirmed_secrets_hex: Vec<String> = self
+            .unconfirmed_secrets
+            .iter()
+            .map(|secret| secret.to_string())
+            .collect();
+        state.serialize_field("unconfirmed_secrets", &unconfirmed_secrets_hex)?;
+
+        state.end()
+    }
+}
+
+impl<'de> Deserialize<'de> for SecretsStore {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        // Deserialize the structure as a map
+        #[derive(Deserialize)]
+        struct SecretsStoreHelper {
+            shared_secrets: HashMap<String, String>,
+            unconfirmed_secrets: Vec<String>,
+        }
+
+        let helper = SecretsStoreHelper::deserialize(deserializer)?;
+
+        let mut shared_secrets = HashMap::with_capacity(helper.shared_secrets.len());
+        for (key_str, value_str) in helper.shared_secrets {
+            let key = SilentPaymentAddress::try_from(key_str).map_err(serde::de::Error::custom)?; // Convert String to SilentPaymentAddress
+            let value = AnkSharedSecretHash::from_str(&value_str).map_err(serde::de::Error::custom)?; // Convert hex string back to Vec<u8>
+
+            shared_secrets.insert(key, value);
+        }
+
+        let mut unconfirmed_secrets = Vec::with_capacity(helper.unconfirmed_secrets.len());
+        for secret_str in helper.unconfirmed_secrets {
+            let secret_bytes = AnkSharedSecretHash::from_str(&secret_str).map_err(serde::de::Error::custom)?;
+            unconfirmed_secrets.push(secret_bytes);
+        }
+
+        Ok(SecretsStore {
+            shared_secrets,
+            unconfirmed_secrets,
+        })
+    }
+}
+
+impl SecretsStore {
+    pub fn new() -> Self {
+        Self { 
+            shared_secrets: HashMap::new(), 
+            unconfirmed_secrets: Vec::new() 
+        }
+    }
+
+    pub fn add_unconfirmed_secret(&mut self, new_secret: AnkSharedSecretHash) {
+        self.unconfirmed_secrets.push(new_secret);
+    }
+
+    pub fn confirm_secret_for_address(&mut self, secret: AnkSharedSecretHash, address: SilentPaymentAddress) {
+        if let Some(pos) = self.unconfirmed_secrets.iter()
+            .position(|s| *s == secret)
+        {
+            self.shared_secrets.insert(address, self.unconfirmed_secrets.swap_remove(pos));
+        } else {
+            // We didn't know about that secret, just add it
+            // TODO if we already had a secret for this address we just replace it for now
+            self.shared_secrets.insert(address, secret);
+        }
+    }
+
+    pub fn get_secret_for_address(&self, address: SilentPaymentAddress) -> Option<&AnkSharedSecretHash> {
+        self.shared_secrets.get(&address)
+    }
+
+    pub fn try_decrypt(&self, cipher: &[u8]) -> Result<(AnkSharedSecretHash, Vec<u8>)> {
+        let nonce = Nonce::from_slice(&cipher[..12]);
+
+        for (address, secret) in self.shared_secrets.iter() {
+            log::debug!("Attempting decryption with key {} for {}", secret, address);
+            let engine = Aes256Gcm::new(secret.as_byte_array().into());
+
+            if let Ok(plain) = engine.decrypt(
+                &nonce,
+                Payload {
+                    msg: &cipher[12..],
+                    aad: AAD,
+                },
+            ) {
+                return Ok((*secret, plain));
+            }
+        }
+
+        // We try unconfirmed secrets
+        for secret in self.unconfirmed_secrets.iter() {
+            log::debug!("Attempting decryption with unconfirmed key {}", secret);
+            let engine = Aes256Gcm::new(secret.as_byte_array().into());
+
+            if let Ok(plain) = engine.decrypt(
+                &nonce,
+                Payload {
+                    msg: &cipher[12..],
+                    aad: AAD,
+                },
+            ) {
+                return Ok((*secret, plain));
+            }
+        }
+
+        Err(Error::msg("Failed to decrypt message"))
+    }
+}