sdk_common/src/crypto.rs

use std::collections::HashMap;

use anyhow::{Error, Result};
use serde::{Deserialize, Serialize};
use sp_client::{
    bitcoin::{
        hex::{DisplayHex, FromHex}, key::constants::SECRET_KEY_SIZE, Txid
    },
    silentpayments::{
        bitcoin_hashes::{sha256t_hash_newtype, Hash, HashEngine},
        utils::SilentPaymentAddress,
        secp256k1::PublicKey
    },
};
use tsify::Tsify;

use aes_gcm::aead::{Aead, Payload};
pub use aes_gcm::{AeadCore, Aes256Gcm, KeyInit};
use rand::thread_rng;

const AAD: &[u8] = "4nk".as_bytes();

const HALFKEYSIZE: usize = SECRET_KEY_SIZE / 2;

const THIRTYTWO: usize = 32;

#[derive(Debug)]
pub struct SharedPoint([u8; 64]);

impl SharedPoint {
    pub fn as_inner(&self) -> &[u8; 64] {
        &self.0
    }
}

#[derive(Debug, Serialize, Deserialize, Tsify, Clone, Default, PartialEq)]
#[tsify(from_wasm_abi, into_wasm_abi)]
pub struct AnkSharedSecret {
    secret: String,
}

impl AnkSharedSecret {
    pub fn new(shared_point: PublicKey) -> Self {
        let mut shared_point_bin = [0u8;64];
        shared_point_bin.copy_from_slice(&shared_point.serialize_uncompressed()[1..]);
        let secret = AnkSharedSecretHash::from_shared_point(shared_point_bin).to_byte_array();
        Self { secret: secret.to_lower_hex_string() }
    }

    pub fn to_byte_array(&self) -> [u8; 32] {
        let bytes = Vec::from_hex(&self.secret).unwrap();
        let mut buf = [0u8;32];
        buf.copy_from_slice(&bytes);
        buf
    }
}

sha256t_hash_newtype! {
    pub struct AnkSharedSecretTag = hash_str("4nk/AnkSharedSecret");

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

impl AnkSharedSecretHash {
    pub fn from_shared_point(shared_point: [u8; 64]) -> Self {
        let mut eng = AnkSharedSecretHash::engine();
        eng.input(&shared_point);
        AnkSharedSecretHash::from_engine(eng)
    }
}

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,
    Arbitrary,
}

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,
        aes_key: [u8;32], 
    ) -> Result<Self> {
        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())
            }
            Purpose::Arbitrary => {
                let arbitrary = self.decrypt_arbitrary()?;
                Ok(arbitrary)
            }
        }
    }

    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)
    }

    fn decrypt_arbitrary(&self) -> Result<Vec<u8>> {
        let cipher = Aes256Gcm::new(&self.aes_key.into());
        let payload = Payload {
            msg: &self.cipher_text,
            aad: AAD,
        };
        let plain = cipher
            .decrypt(&self.nonce.into(), payload)
            .map_err(|e| Error::msg(format!("{}", e)))?;
        Ok(plain)
    }
}

pub struct Aes256Encryption {
    pub purpose: Purpose,
    plaintext: Vec<u8>,
    aes_key: [u8; 32],
    nonce: [u8; 12],
    shared_secrets: HashMap<Txid, HashMap<SilentPaymentAddress, AnkSharedSecret>>,
}

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, AnkSharedSecret>>,
    ) {
        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.to_byte_array())
                    .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(),
            Purpose::Arbitrary => self.encrypt_arbitrary(),
        }
    }

    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)
    }

    fn encrypt_arbitrary(&self) -> Result<CipherText> {
        let cipher = Aes256Gcm::new(&self.aes_key.into());
        let payload = Payload {
            msg: &self.plaintext,
            aad: AAD,
        };
        let cipher_text = cipher
            .encrypt(&self.nonce.into(), payload)
            .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::{io::Read, str::FromStr};

    use sp_client::bitcoin::hex::FromHex;

    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);

        assert!(aes_dec.is_ok());
    }
}