sdk_common/src/secrets.rs

use anyhow::{Result, Error};
use tsify::Tsify;
use crate::aes_gcm::aead::{Aead, Payload};
use crate::aes_gcm::Nonce;
use crate::sp_client::bitcoin::hashes::Hash;
use crate::sp_client::silentpayments::SilentPaymentAddress;
use crate::crypto::{Aes256Gcm, AnkSharedSecretHash, KeyInit, AAD};
use serde::ser::SerializeStruct;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::collections::HashMap;
use std::str::FromStr;

#[derive(Debug, Clone, Default, PartialEq, Tsify)]
#[tsify(into_wasm_abi)]
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);
    }

    /// Returns the previous secret for this address, if any
    pub fn confirm_secret_for_address(&mut self, secret: AnkSharedSecretHash, address: SilentPaymentAddress) -> Option<AnkSharedSecretHash> {
        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 get_all_confirmed_secrets(&self) -> HashMap<SilentPaymentAddress, AnkSharedSecretHash> {
        self.shared_secrets.clone()
    }

    pub fn get_all_unconfirmed_secrets(&self) -> Vec<AnkSharedSecretHash> {
        self.unconfirmed_secrets.clone()
    }

    pub fn remove_secret_for_address(&mut self, address: SilentPaymentAddress) -> Result<(SilentPaymentAddress, AnkSharedSecretHash)> {
        if let Some(removed_secret) = self.shared_secrets.remove(&address) {
            return Ok((address, removed_secret));
        } else {
            return Err(Error::msg("Secret doesn't exist"));
        }
    }

    pub fn remove_unconfirmed_secret(&mut self, secret: AnkSharedSecretHash) -> Result<()> {
        if let Some(i) = self.unconfirmed_secrets.iter().position(|s| *s == secret) {
            self.unconfirmed_secrets.swap_remove(i);
            Ok(())
        } else {
            Err(Error::msg("Unknown secret"))
        }
    }

    pub fn try_decrypt(&self, cipher: &[u8]) -> Result<(AnkSharedSecretHash, Vec<u8>)> {
        let nonce = Nonce::from_slice(&cipher[..12]);

        for (_, secret) in self.shared_secrets.iter() {
            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() {
            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"))
    }
}