From dd797bd981ad52b7099f303352b0fecebfab649d Mon Sep 17 00:00:00 2001
From: Sosthene <contact@sosthene.net>
Date: Thu, 15 Aug 2024 16:15:26 +0200
Subject: [PATCH] Update Pcd/Prd

---
 src/network.rs | 66 +++++++++++++++++++++++++++++++++++---------------
 1 file changed, 47 insertions(+), 19 deletions(-)

diff --git a/src/network.rs b/src/network.rs
index a17d75d..9de12d4 100644
--- a/src/network.rs
+++ b/src/network.rs
@@ -1,13 +1,16 @@
-use std::default;
+use std::{default, fmt};
 use std::str::FromStr;
 
-use aes_gcm::{AeadCore, Aes256Gcm};
+use aes_gcm::{AeadCore, Aes256Gcm, KeyInit};
 use anyhow::{Error, Result};
 use js_sys::Date;
 use rand::{thread_rng, RngCore};
-use serde::{Deserialize, Serialize};
+use serde::de::{MapAccess, Visitor};
+use serde::ser::SerializeStruct;
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use sp_client::bitcoin::consensus::serialize;
-use sp_client::bitcoin::hashes::Hash;
+use sp_client::bitcoin::hashes::sha256::Hash;
+// use sp_client::bitcoin::hashes::Hash;
 use sp_client::bitcoin::hex::{DisplayHex, FromHex};
 use sp_client::bitcoin::secp256k1::PublicKey;
 use sp_client::bitcoin::{BlockHash, OutPoint, Transaction};
@@ -109,28 +112,42 @@ impl NewTxMessage {
     }
 }
 
-#[derive(Debug, Serialize, Deserialize, Clone, Tsify)]
+#[derive(Debug, Clone, Serialize, Deserialize, Tsify)]
 #[tsify(into_wasm_abi, from_wasm_abi)]
+#[allow(non_camel_case_types)]
 pub struct Prd {
     pub sender: String,
-    pub keys: Vec<String>,
-    pub pcd_commitment: String, // hash of the pcd
+    pub key: [u8;32],
+    pub pcd_commitment: Hash, // hash of the pcd
     pub error: Option<AnkError>,
 }
 
 impl Prd {
-    pub fn new(sender: String, keys: Vec<String>, pcd_commitment: String) -> Self {
+    pub fn new(sender: SilentPaymentAddress, key: [u8; 32], pcd_commitment: Hash) -> Self {
         Self {
-            sender,
-            keys,
+            sender: sender.into(),
+            key,
             pcd_commitment,
             error: None,
         }
     }
+
+    pub fn encrypt_pcd(&self, pcd: &Pcd) -> Result<Vec<u8>> {
+        let plain = pcd.to_string().into_bytes();
+        let nonce: [u8; 12] = Aes256Gcm::generate_nonce(thread_rng()).into();
+        let aes_encrypt = Aes256Encryption::import_key(Purpose::Arbitrary, plain, self.key, nonce)?;
+        let cipher = aes_encrypt.encrypt_with_aes_key()?;
+        Ok(cipher)
+    }
+
+    pub fn to_string(&self) -> String {
+        serde_json::to_string(self).unwrap()
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, Tsify)]
 #[tsify(into_wasm_abi, from_wasm_abi)]
+#[allow(non_camel_case_types)]
 pub struct Pcd {
     pub message: String,
     pub error: Option<AnkError>,
@@ -149,8 +166,7 @@ impl Pcd {
     }
 }
 
-#[derive(Debug, Serialize, Deserialize, Tsify)]
-#[tsify(into_wasm_abi, from_wasm_abi)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct Envelope {
     pub flag: AnkFlag,
     pub content: String,
@@ -182,7 +198,8 @@ pub enum CachedMessageStatus {
     Pairing,
     Login,
     CipherWaitingTx,
-    TxWaitingCipher,
+    TxWaitingPrd,
+    GotPrdWaitingPcd,
     SentWaitingConfirmation, // we're sender and wait for commited_in to be spent
     ReceivedMustConfirm,     // we're receiver and we spend commited_in to sender
     MustSpendConfirmation,   // we're sender and we must spend confirmed_by
@@ -208,15 +225,12 @@ pub enum CachedMessageStatus {
 pub struct CachedMessage {
     pub id: u32,
     pub status: CachedMessageStatus,
-    pub ciphertext: Option<String>, // Needed when we are waiting for a key in transaction, discarded after
-    pub plaintext: Vec<String>,     // Append new messages received while in Trusted state
     pub commited_in: Option<OutPoint>,
     pub tied_by: Option<u32>, // index of the output that ties the proposal
     pub commitment: Option<String>, // content of the op_return
     pub sender: Option<String>, // Never None when message sent
     pub recipient: Option<String>, // Never None when message sent
     pub shared_secret: Option<String>, // Never None when message sent
-    pub key: Option<String>,   // Never None when message sent
     pub prd_cipher: Option<String>, // When we receive message we can't decrypt we only have this and commited_in_tx
     pub prd: Option<Prd>, // Never None when message sent
     pub pcd_cipher: Option<String>, 
@@ -246,7 +260,7 @@ impl CachedMessage {
         serde_json::to_string(self).unwrap()
     }
 
-    pub fn try_decrypt_cipher(&self, cipher: Vec<u8>) -> Result<Vec<u8>> {
+    pub fn try_decrypt_prd(&self, cipher: Vec<u8>) -> Result<Vec<u8>> {
         if self.shared_secret.is_none() {
             return Err(Error::msg(
                 "Can't try decrypt this message, there's no shared secret",
@@ -259,15 +273,29 @@ impl CachedMessage {
         aes_decrypt.decrypt_with_key()
     }
 
+    pub fn try_decrypt_pcd(&self, cipher: Vec<u8>) -> Result<Vec<u8>> {
+        if self.prd.is_none() {
+            return Err(Error::msg(
+                "Can't try decrypt this message, there's no prd",
+            ));
+        }
+        let aes_decrypt = Aes256Decryption::new(Purpose::Arbitrary, cipher, self.prd.as_ref().unwrap().key)?;
+
+        aes_decrypt.decrypt_with_key()
+    }
+
     pub fn try_decrypt_with_shared_secret(&self, shared_secret: [u8; 32]) -> Result<Vec<u8>> {
         if self.prd_cipher.is_none() || self.shared_secret.is_some() {
             return Err(Error::msg(
                 "Can't try decrypt this message, ciphertext is none or shared_secret already found",
             ));
         }
-        let cipher_bin = Vec::from_hex(self.ciphertext.as_ref().unwrap())?;
-        let aes_decrypt = Aes256Decryption::new(Purpose::Arbitrary, cipher_bin, shared_secret)?;
 
+        let mut key = [0u8; 32];
+        key.copy_from_slice(&shared_secret);
+        let cipher = Vec::from_hex(&self.prd_cipher.as_ref().unwrap()).expect("Shouldn't keep an invalid hex as cipher");
+
+        let aes_decrypt = Aes256Decryption::new(Purpose::Arbitrary, cipher, key)?;
         aes_decrypt.decrypt_with_key()
     }
 }