sdk_common/src/prd.rs

use std::collections::{HashMap, HashSet};

use anyhow::Result;
use serde::{Deserialize, Serialize};

use serde_json::{Map, Value};
use sp_client::bitcoin::hashes::{sha256t_hash_newtype, Hash, HashEngine};
use sp_client::bitcoin::secp256k1::{PublicKey, SecretKey};
use sp_client::bitcoin::OutPoint;
use sp_client::silentpayments::utils::SilentPaymentAddress;
use sp_client::spclient::SpWallet;
use tsify::Tsify;

use crate::pcd::{Member, RoleDefinition};
use crate::signature::{AnkHash, AnkMessageHash, Proof};

#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
#[allow(non_camel_case_types)]
pub enum PrdType {
    #[default]
    None,
    Connect,
    Message,
    Update, // Update an existing process
    List,   // request a list of items
    Response, // Validate (or disagree) with a prd update
    Confirm, // Confirm we received an update
    TxProposal, // Send a psbt asking for recipient signature
}

sha256t_hash_newtype! {
    pub struct AnkPrdTag = hash_str("4nk/Prd");

    #[hash_newtype(forward)]
    pub struct AnkPrdHash(_);
}

impl AnkPrdHash {
    pub fn from_value(value: &Value) -> Self {
        let mut eng = AnkPrdHash::engine();
        eng.input(value.to_string().as_bytes());
        AnkPrdHash::from_engine(eng)
    }

    pub fn from_map(map: &Map<String, Value>) -> Self {
        let value = Value::Object(map.clone());
        let mut eng = AnkPrdHash::engine();
        eng.input(value.to_string().as_bytes());
        AnkPrdHash::from_engine(eng)
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
#[allow(non_camel_case_types)]
pub struct Prd {
    pub prd_type: PrdType,
    pub root_commitment: String,
    pub sender: String,
    pub keys: Map<String, Value>, // key is a key in pcd, value is the key to decrypt it
    pub pcd_commitments: Value,
    pub validation_tokens: Vec<Proof>,
    pub payload: String,      // Additional information depending on the type
    pub proof: Option<Proof>, // This must be None up to the creation of the network message
}

impl Prd {
    /// We answer to ack we received a transaction and got the shared_secret
    /// If validation_tokens is empty we put the proof into it and return it
    /// If validation_tokens contains a valid proof signed by ourselves of empty prd, 
    /// we confirm the secret if necessary and don't return anything
    pub fn new_connect(sender: Member, secret_hash: AnkMessageHash, previous_proof: Option<Proof>) -> Self {
        let validation_tokens = if let Some(proof) = previous_proof { vec![proof] } else { vec![] };
        Self {
            prd_type: PrdType::Connect,
            root_commitment: String::default(),
            pcd_commitments: Value::Null,
            sender: serde_json::to_string(&sender).unwrap(),
            validation_tokens,
            keys: Map::new(),
            payload: secret_hash.to_string(),
            proof: None,
        }
    }

    pub fn new_update(
        root_commitment: OutPoint,
        sender: String, // Should take Member as argument
        roles: HashMap<String, RoleDefinition>,
        keys: Map<String, Value>,
        pcd_commitments: Value,
    ) -> Self {
        Self {
            prd_type: PrdType::Update,
            root_commitment: root_commitment.to_string(),
            sender,
            validation_tokens: vec![],
            keys,
            pcd_commitments,
            payload: serde_json::to_string(&roles).expect("We're confident it's serializable"),
            proof: None,
        }
    }

    pub fn new_response(
        root_commitment: OutPoint,
        sender: String,
        validation_tokens: Vec<Proof>,
        pcd_commitments: Value,
    ) -> Self {
        Self {
            prd_type: PrdType::Response,
            root_commitment: root_commitment.to_string(),
            sender,
            validation_tokens: validation_tokens,
            keys: Map::new(),
            pcd_commitments,
            payload: String::default(),
            proof: None,
        }
    }

    pub fn new_confirm(
        root_commitment: OutPoint,
        sender: Member,
        pcd_commitments: Value,
    ) -> Self {
        Self {
            prd_type: PrdType::Confirm,
            root_commitment: root_commitment.to_string(),
            pcd_commitments,
            sender: serde_json::to_string(&sender).unwrap(),
            validation_tokens: vec![],
            keys: Map::new(),
            payload: String::default(),
            proof: None,
        }
    }

    pub fn extract_from_message(plain: &[u8], local_address: SilentPaymentAddress) -> Result<Self> {
        let prd: Prd = serde_json::from_slice(plain)?;

        // check that the proof is consistent
        if let Some(proof) = prd.proof {
            let proof_key = proof.get_key();
            let local_spend_key = local_address.get_spend_key();
            // If it's our own device key we abort
            if proof_key == local_spend_key {
                return Err(anyhow::Error::msg("Proof signed by ourselves, we are parsing our own message"));
            }
            // take the spending keys in sender
            let sender: Member = serde_json::from_str(&prd.sender)?;
            let addresses = sender.get_addresses();
            let mut spend_keys: Vec<PublicKey> = vec![];
            for address in addresses {
                spend_keys.push(
                    <SilentPaymentAddress>::try_from(address)?
                        .get_spend_key()
                );
            }
            // The key in proof must be one of the sender keys
            let mut known_key = false;
            for key in spend_keys {
                if key == proof_key {
                    known_key = true;
                    break;
                }
            }
            if !known_key {
                return Err(anyhow::Error::msg("Proof signed with an unknown key"));
            }
            proof.verify()?;
        } else {
            log::warn!("No proof for prd with root_commitment {}", prd.root_commitment);
        }
        Ok(prd)
    }

    pub fn filter_keys(&mut self, to_keep: HashSet<String>) {
        let current_keys = self.keys.clone();
        let filtered_keys: Map<String, Value> = current_keys
            .into_iter()
            .filter(|(field, _)| to_keep.contains(field))
            .collect();
        self.keys = filtered_keys;
    }

    /// Generate the signed proof and serialize to send over the network
    pub fn to_network_msg(&self, sp_wallet: &SpWallet) -> Result<String> {
        let spend_sk: SecretKey = sp_wallet.get_client().get_spend_key().try_into()?;
        let mut to_sign = self.clone(); // we sign the whole prd, incl the keys, for each recipient

        let message_hash =
            AnkHash::Message(AnkMessageHash::from_message(to_sign.to_string().as_bytes()));

        let proof = Proof::new(message_hash, spend_sk);

        to_sign.proof = Some(proof);

        Ok(to_sign.to_string())
    }

    pub fn to_string(&self) -> String {
        serde_json::to_string(self).unwrap()
    }

    pub fn to_value(&self) -> Value {
        serde_json::to_value(self).unwrap()
    }
}