From 3dff7922187b8647775641d9d9ba1fd1b06ff3c5 Mon Sep 17 00:00:00 2001
From: Sosthene00 <674694@protonmail.ch>
Date: Tue, 9 Apr 2024 12:58:41 +0200
Subject: [PATCH] Add crypto mod

---
 Cargo.toml    |   5 +-
 src/crypto.rs | 441 ++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs    |   4 +-
 3 files changed, 445 insertions(+), 5 deletions(-)
 create mode 100644 src/crypto.rs

diff --git a/Cargo.toml b/Cargo.toml
index 6dcc1a1..1737a6a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,9 +7,10 @@ edition = "2021"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
+anyhow = "1.0"
 serde = { version = "1.0.193", features = ["derive"] }
 serde_json = "1.0.108"
-sp_backend = { git = "https://github.com/Sosthene00/sp-backend.git", branch = "master" }
+sp_backend = { git = "https://github.com/Sosthene00/sp-backend.git", branch = "sp_client" }
 uuid = { version = "1.6.1", features = ["serde", "v4"] }
 aes-gcm = "0.10.3"
-aes = "0.8.3"
+rand = "0.8.5"
diff --git a/src/crypto.rs b/src/crypto.rs
new file mode 100644
index 0000000..3873b1d
--- /dev/null
+++ b/src/crypto.rs
@@ -0,0 +1,441 @@
+use std::collections::HashMap;
+
+use anyhow::{Error, Result};
+use sp_backend::{
+    bitcoin::{
+        consensus::serde::hex,
+        hex::DisplayHex,
+        key::constants::SECRET_KEY_SIZE,
+        secp256k1::{ecdh::SharedSecret, SecretKey},
+        Txid,
+    },
+    silentpayments::sending::SilentPaymentAddress,
+};
+
+use serde::{Deserialize, Serialize};
+use serde_json::{json, Value};
+
+use aes_gcm::{
+    aead::{Aead, AeadInPlace, KeyInit, Nonce},
+    AeadCore, Aes256Gcm, AesGcm, Key, TagSize,
+    aes::{Aes256, cipher::{generic_array::GenericArray, consts::{U32, U8}}},
+};
+use rand::thread_rng;
+
+const HALFKEYSIZE: usize = SECRET_KEY_SIZE / 2;
+
+const THIRTYTWO: usize = 32;
+
+pub struct HalfKey([u8; HALFKEYSIZE]);
+
+impl TryFrom<Vec<u8>> for HalfKey {
+    type Error = anyhow::Error;
+    fn try_from(value: Vec<u8>) -> std::prelude::v1::Result<Self, Error> {
+        if value.len() == HALFKEYSIZE {
+            let mut buf = [0u8; HALFKEYSIZE];
+            buf.copy_from_slice(&value);
+            Ok(HalfKey(buf))
+        } else {
+            Err(Error::msg("Invalid length for HalfKey"))
+        }
+    }
+}
+
+impl HalfKey {
+    pub fn as_slice(&self) -> &[u8] {
+        &self.0
+    }
+
+    pub fn to_inner(&self) -> Vec<u8> {
+        self.0.to_vec()
+    }
+}
+
+pub enum Purpose {
+    Login,
+    ThirtyTwoBytes,
+}
+
+pub type CipherText = Vec<u8>;
+
+pub type EncryptedKey = Vec<u8>;
+
+pub struct Aes256Decryption {
+    pub purpose: Purpose,
+    cipher_text: CipherText,
+    aes_key: [u8; 32],
+    nonce: [u8; 12],
+}
+
+impl Aes256Decryption {
+    pub fn new(
+        purpose: Purpose,
+        cipher_text: CipherText,
+        encrypted_aes_key: Vec<u8>, // If shared_secret is none this is actually the aes_key
+        shared_secret: Option<SharedSecret>, // We don't need that for certain purpose, like Login
+    ) -> Result<Self> {
+        let mut aes_key = [0u8; 32];
+        if let Some(shared_secret) = shared_secret {
+            if encrypted_aes_key.len() <= 12 {
+                return Err(Error::msg("encrypted_aes_key is shorter than nonce length"));
+            } // Actually we could probably test that if the remnant is not a multiple of 32, something's wrong
+              // take the first 12 bytes form encrypted_aes_key as nonce
+            let (decrypt_key_nonce, encrypted_key) = encrypted_aes_key.split_at(12);
+            // decrypt key with shared_secret obtained from transaction
+            let decrypt_key_cipher = Aes256Gcm::new_from_slice(shared_secret.as_ref())
+                .map_err(|e| Error::msg(format!("{}", e)))?;
+            let aes_key_plain = decrypt_key_cipher
+                .decrypt(decrypt_key_nonce.into(), encrypted_key)
+                .map_err(|e| Error::msg(format!("{}", e)))?;
+            if aes_key_plain.len() != 32 {
+                return Err(Error::msg("Invalid length for decrypted key"));
+            }
+            aes_key.copy_from_slice(&aes_key_plain);
+        } else {
+            if encrypted_aes_key.len() != 32 {
+                return Err(Error::msg("Invalid length for decrypted key"));
+            }
+            aes_key.copy_from_slice(&encrypted_aes_key);
+        }
+        if cipher_text.len() <= 12 {
+            return Err(Error::msg("cipher_text is shorter than nonce length"));
+        }
+        let (message_nonce, message_cipher) = cipher_text.split_at(12);
+        let mut nonce = [0u8; 12];
+        nonce.copy_from_slice(message_nonce);
+        Ok(Self {
+            purpose,
+            cipher_text: message_cipher.to_vec(),
+            aes_key,
+            nonce,
+        })
+    }
+
+    pub fn decrypt_with_key(&self) -> Result<Vec<u8>> {
+        match self.purpose {
+            Purpose::Login => {
+                let half_key = self.decrypt_login()?;
+                Ok(half_key.to_inner())
+            }
+            Purpose::ThirtyTwoBytes => {
+                let thirty_two_buf = self.decrypt_thirty_two()?;
+                Ok(thirty_two_buf.to_vec())
+            }
+        }
+    }
+
+    fn decrypt_login(&self) -> Result<HalfKey> {
+        let cipher = Aes256Gcm::new(&self.aes_key.into());
+        let plain = cipher
+            .decrypt(&self.nonce.into(), &*self.cipher_text)
+            .map_err(|e| Error::msg(format!("{}", e)))?;
+        if plain.len() != SECRET_KEY_SIZE / 2 {
+            return Err(Error::msg("Plain text of invalid lenght for a login"));
+        }
+        let mut key_half = [0u8; SECRET_KEY_SIZE / 2];
+        key_half.copy_from_slice(&plain);
+        Ok(HalfKey(key_half))
+    }
+
+    fn decrypt_thirty_two(&self) -> Result<[u8; THIRTYTWO]> {
+        let cipher = Aes256Gcm::new(&self.aes_key.into());
+        let plain = cipher
+            .decrypt(&self.nonce.into(), &*self.cipher_text)
+            .map_err(|e| Error::msg(format!("{}", e)))?;
+        if plain.len() != THIRTYTWO {
+            return Err(Error::msg("Plain text of invalid length, should be 32"));
+        }
+        let mut thirty_two = [0u8; THIRTYTWO];
+        thirty_two.copy_from_slice(&plain);
+        Ok(thirty_two)
+    }
+}
+
+pub struct Aes256Encryption {
+    pub purpose: Purpose,
+    plaintext: Vec<u8>,
+    aes_key: [u8; 32],
+    nonce: [u8; 12],
+    shared_secrets: HashMap<Txid, HashMap<SilentPaymentAddress, SharedSecret>>,
+}
+
+impl Aes256Encryption {
+    pub fn new(purpose: Purpose, plaintext: Vec<u8>) -> Result<Self> {
+        let mut rng = thread_rng();
+        let aes_key: [u8; 32] = Aes256Gcm::generate_key(&mut rng).into();
+        let nonce: [u8; 12] = Aes256Gcm::generate_nonce(&mut rng).into();
+        Self::import_key(purpose, plaintext, aes_key, nonce)
+    }
+
+    pub fn set_shared_secret(
+        &mut self,
+        shared_secrets: HashMap<Txid, HashMap<SilentPaymentAddress, SharedSecret>>,
+    ) {
+        self.shared_secrets = shared_secrets;
+    }
+
+    pub fn encrypt_keys_with_shared_secrets(
+        &self,
+    ) -> Result<HashMap<SilentPaymentAddress, EncryptedKey>> {
+        let mut res = HashMap::new();
+        let mut rng = thread_rng();
+
+        for (_, sp_address2shared_secret) in self.shared_secrets.iter() {
+            for (sp_address, shared_secret) in sp_address2shared_secret {
+                let cipher = Aes256Gcm::new_from_slice(shared_secret.as_ref())
+                    .map_err(|e| Error::msg(format!("{}", e)))?;
+                let nonce = Aes256Gcm::generate_nonce(&mut rng);
+                let encrypted_key = cipher
+                    .encrypt(&nonce, self.aes_key.as_slice())
+                    .map_err(|e| Error::msg(format!("{}", e)))?;
+
+                let mut ciphertext = Vec::<u8>::with_capacity(nonce.len() + encrypted_key.len());
+                ciphertext.extend(nonce);
+                ciphertext.extend(encrypted_key);
+
+                res.insert(sp_address.to_owned(), ciphertext);
+            }
+        }
+        Ok(res)
+    }
+
+    pub fn import_key(
+        purpose: Purpose,
+        plaintext: Vec<u8>,
+        aes_key: [u8; 32],
+        nonce: [u8; 12],
+    ) -> Result<Self> {
+        if plaintext.len() == 0 {
+            return Err(Error::msg("Can't create encryption for an empty message"));
+        }
+        Ok(Self {
+            purpose,
+            plaintext,
+            aes_key,
+            nonce,
+            shared_secrets: HashMap::new(),
+        })
+    }
+
+    pub fn encrypt_with_aes_key(&self) -> Result<CipherText> {
+        match self.purpose {
+            Purpose::Login => self.encrypt_login(),
+            Purpose::ThirtyTwoBytes => self.encrypt_thirty_two(),
+        }
+    }
+
+    fn encrypt_login(&self) -> Result<CipherText> {
+        let half_key: HalfKey = self.plaintext.clone().try_into()?;
+        let cipher = Aes256Gcm::new(&self.aes_key.into());
+        let cipher_text = cipher
+            .encrypt(&self.nonce.into(), half_key.as_slice())
+            .map_err(|e| Error::msg(format!("{}", e)))?;
+        let mut res = Vec::with_capacity(self.nonce.len() + cipher_text.len());
+        res.extend_from_slice(&self.nonce);
+        res.extend_from_slice(&cipher_text);
+        Ok(res)
+    }
+
+    fn encrypt_thirty_two(&self) -> Result<CipherText> {
+        if self.plaintext.len() != 32 {
+            return Err(Error::msg("Invalid length, should be 32"));
+        }
+        let mut thirty_two = [0u8; 32];
+        thirty_two.copy_from_slice(&self.plaintext);
+        let cipher = Aes256Gcm::new(&self.aes_key.into());
+        let cipher_text = cipher
+            .encrypt(&self.nonce.into(), thirty_two.as_slice())
+            .map_err(|e| Error::msg(format!("{}", e)))?;
+        let mut res = Vec::with_capacity(self.nonce.len() + cipher_text.len());
+        res.extend_from_slice(&self.nonce);
+        res.extend_from_slice(&cipher_text);
+        Ok(res)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    use super::*;
+
+    const ALICE_SP_ADDRESS: &str = "tsp1qqw3lqr6xravz9nf8ntazgwwl0fqv47kfjdxsnxs6eutavqfwyv5q6qk97mmyf6dtkdyzqlu2zv6h9j2ggclk7vn705q5u2phglpq7yw3dg5rwpdz";
+    const BOB_SP_ADDRESS: &str = "tsp1qq2hlsgrj0gz8kcfkf9flqw5llz0u2vr04telqndku9mcqm6dl4fhvq60t8r78srrf56w9yr7w9e9dusc2wjqc30up6fjwnh9mw3e3veqegdmtf08";
+    const TRANSACTION: &str = "4e6d03dec558e1b6624f813bf2da7cd8d8fb1c2296684c08cf38724dcfd8d10b";
+    const ALICE_SHARED_SECRET: &str =
+        "ccf02d364c2641ca129a3fdf49de57b705896e233f7ba6d738991993ea7e2106";
+    const BOB_SHARED_SECRET: &str =
+        "15ef3e377fb842e81de52dbaaea8ba30aeb051a81043ee19264afd27353da521";
+
+    #[test]
+    fn new_aes_empty_plaintext() {
+        let plaintext = Vec::new();
+        let aes_enc = Aes256Encryption::new(Purpose::Login, plaintext);
+
+        assert!(aes_enc.is_err());
+    }
+
+    #[test]
+    fn aes_encrypt_login_invalid_length() {
+        let plaintext = "example";
+        let aes_enc_short = Aes256Encryption::new(Purpose::Login, plaintext.as_bytes().to_vec());
+
+        assert!(aes_enc_short.is_ok());
+
+        let cipher = aes_enc_short.unwrap().encrypt_with_aes_key();
+
+        assert!(cipher.is_err());
+
+        let plaintext = [1u8; 64];
+        let aes_enc_long = Aes256Encryption::new(Purpose::Login, plaintext.to_vec());
+
+        assert!(aes_enc_long.is_ok());
+
+        let cipher = aes_enc_long.unwrap().encrypt_with_aes_key();
+
+        assert!(cipher.is_err());
+    }
+
+    #[test]
+    fn aes_encrypt_login() {
+        let plaintext = [1u8; HALFKEYSIZE];
+        let aes_key = Aes256Gcm::generate_key(&mut thread_rng());
+        let nonce = Aes256Gcm::generate_nonce(&mut thread_rng());
+        let aes_enc = Aes256Encryption::import_key(
+            Purpose::Login,
+            plaintext.to_vec(),
+            aes_key.into(),
+            nonce.into(),
+        );
+
+        assert!(aes_enc.is_ok());
+
+        let cipher = aes_enc.unwrap().encrypt_with_aes_key();
+
+        assert!(cipher.is_ok());
+
+        let mut plain_key = [0u8; 32];
+        plain_key.copy_from_slice(&aes_key.to_vec());
+
+        let aes_dec =
+            Aes256Decryption::new(Purpose::Login, cipher.unwrap(), plain_key.to_vec(), None);
+
+        assert!(aes_dec.is_ok());
+    }
+
+    #[test]
+    fn aes_encrypt_key() {
+        let plaintext = [1u8; HALFKEYSIZE];
+        let mut aes_enc = Aes256Encryption::new(Purpose::Login, plaintext.to_vec()).unwrap();
+
+        let mut shared_secrets: HashMap<Txid, _> = HashMap::new();
+        let mut sp_address2shared_secrets: HashMap<SilentPaymentAddress, SharedSecret> =
+            HashMap::new();
+        sp_address2shared_secrets.insert(
+            ALICE_SP_ADDRESS.try_into().unwrap(),
+            SharedSecret::from_str(ALICE_SHARED_SECRET).unwrap(),
+        );
+        shared_secrets.insert(
+            Txid::from_str(TRANSACTION).unwrap(),
+            sp_address2shared_secrets,
+        );
+
+        aes_enc.set_shared_secret(shared_secrets);
+
+        let sp_address2encrypted_keys = aes_enc.encrypt_keys_with_shared_secrets();
+
+        assert!(sp_address2encrypted_keys.is_ok());
+
+        let encrypted_key = sp_address2encrypted_keys
+            .unwrap()
+            .get(&ALICE_SP_ADDRESS.try_into().unwrap())
+            .cloned();
+
+        let ciphertext = aes_enc.encrypt_with_aes_key();
+
+        assert!(ciphertext.is_ok());
+
+        let aes_dec = Aes256Decryption::new(
+            Purpose::Login,
+            ciphertext.unwrap(),
+            encrypted_key.unwrap(),
+            Some(SharedSecret::from_str(ALICE_SHARED_SECRET).unwrap()),
+        );
+
+        assert!(aes_dec.is_ok());
+
+        let retrieved_plain = aes_dec.unwrap().decrypt_with_key();
+
+        assert!(retrieved_plain.is_ok());
+
+        assert!(retrieved_plain.unwrap() == plaintext);
+    }
+
+    #[test]
+    fn aes_encrypt_key_many() {
+        let plaintext = [1u8; THIRTYTWO];
+        let mut aes_enc =
+            Aes256Encryption::new(Purpose::ThirtyTwoBytes, plaintext.to_vec()).unwrap();
+
+        let mut shared_secrets: HashMap<Txid, _> = HashMap::new();
+        let mut sp_address2shared_secrets: HashMap<SilentPaymentAddress, SharedSecret> =
+            HashMap::new();
+        sp_address2shared_secrets.insert(
+            ALICE_SP_ADDRESS.try_into().unwrap(),
+            SharedSecret::from_str(ALICE_SHARED_SECRET).unwrap(),
+        );
+        sp_address2shared_secrets.insert(
+            BOB_SP_ADDRESS.try_into().unwrap(),
+            SharedSecret::from_str(BOB_SHARED_SECRET).unwrap(),
+        );
+        shared_secrets.insert(
+            Txid::from_str(TRANSACTION).unwrap(),
+            sp_address2shared_secrets,
+        );
+
+        aes_enc.set_shared_secret(shared_secrets);
+
+        let mut sp_address2encrypted_keys = aes_enc.encrypt_keys_with_shared_secrets();
+
+        assert!(sp_address2encrypted_keys.is_ok());
+
+        // Alice
+        let encrypted_key = sp_address2encrypted_keys
+            .as_mut()
+            .unwrap()
+            .get(&ALICE_SP_ADDRESS.try_into().unwrap())
+            .cloned();
+
+        let ciphertext = aes_enc.encrypt_with_aes_key();
+
+        let aes_dec = Aes256Decryption::new(
+            Purpose::ThirtyTwoBytes,
+            ciphertext.unwrap(),
+            encrypted_key.unwrap(),
+            Some(SharedSecret::from_str(ALICE_SHARED_SECRET).unwrap()),
+        );
+
+        let retrieved_plain = aes_dec.unwrap().decrypt_with_key();
+
+        assert!(retrieved_plain.unwrap() == plaintext);
+
+        // Bob
+        let encrypted_key = sp_address2encrypted_keys
+            .unwrap()
+            .get(&BOB_SP_ADDRESS.try_into().unwrap())
+            .cloned();
+
+        let ciphertext = aes_enc.encrypt_with_aes_key();
+
+        let aes_dec = Aes256Decryption::new(
+            Purpose::ThirtyTwoBytes,
+            ciphertext.unwrap(),
+            encrypted_key.unwrap(),
+            Some(SharedSecret::from_str(BOB_SHARED_SECRET).unwrap()),
+        );
+
+        let retrieved_plain = aes_dec.unwrap().decrypt_with_key();
+
+        assert!(retrieved_plain.unwrap() == plaintext);
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index c21f7c9..274f0ed 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,3 +1 @@
-pub mod models;
-pub mod wallet;
-pub mod workflows;
+pub mod crypto;