Try moving from PrivateKey to Static Secret

2 files changed, 239 insertions, 28 deletions

diff --git a/src/crypto.rs b/src/crypto.rs
index d967f50..63532ab 100644
--- a/src/crypto.rs
+++ b/src/crypto.rs
@@ -1,22 +1,30 @@
 use std::str::FromStr;
 use tokio_util::bytes::{Buf, BufMut, Bytes, BytesMut};
-use hex::{decode, decode_to_slice};
+use hex::{decode, decode_to_slice, encode};
 
 use ed25519_dalek::hazmat::ExpandedSecretKey;
-use curve25519_dalek::{constants, edwards::CompressedEdwardsY};
+use curve25519_dalek::{Scalar, constants, edwards::CompressedEdwardsY};
 use x25519_dalek::StaticSecret;
-// use aes::{Aes256, cipher::{Block, BlockDecryptMut, KeyInit}};
-
 
+// This seems to be an absolute nightmare.  GenericArray sucks
+// but I can't seem to figure out how to pull it out of this
+// stack of software
+#[allow(deprecated)]
+use aes::cipher::{
+    BlockCipher, BlockEncrypt, BlockDecrypt,
+    generic_array::GenericArray,
+};
+use aes::Aes256;
 use sha2::{Sha256};
 use hmac::{Hmac, Mac};
 type HmacSha256 = Hmac<Sha256>;
 
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub enum MeshcoreCryptoError {
     KeyLengthError,
-    TryFromSliceError
+    TryFromSliceError,
+    HexDecodeError
 }
 
 #[derive(PartialEq, Debug)]
@@ -66,14 +74,79 @@ impl SharedSecret {
         }
     }
 
-    pub fn decrypt(&self, ciphertext: Bytes) -> Bytes {
-        // let mut aes = Aes256::new((&self.0).into());
-
-        // let mut block = Block::clone_from_slice(&ciphertext);
-        // aes.decrypt_block_mut(&mut block);
+    pub fn decrypt(&self, ciphertext: Bytes) -> Result<Bytes, MeshcoreCryptoError> {
+        use aes::cipher::KeyInit;
+
+        if let Ok(mut aes) = Aes256::new_from_slice(&self.0) {
+            let mut text = BytesMut::from(ciphertext);
+
+            // The decryption function works on a 16-byte block of data.
+            // we need to break the input ciphertext into blocks of this size
+            // and work with them individually.  Also, any "short" blocks have
+            // to be zero-padded.
+            // Copy the original length so we can truncate back down
+            let chunk_size = 16;
+            let length = text.len();
+            let remainder = length % chunk_size;
+            // Pad the bytes to it's an even multiple of the chunk size
+
+            if remainder != 0 {
+                let padding = chunk_size - remainder;
+                text.reserve(padding);
+                text.put_bytes(0, padding);
+            }
+            
+            let chunks = text.chunks_exact_mut(chunk_size);
+            for chunk in chunks {
+                #[allow(deprecated)]
+                let mut block = *GenericArray::from_slice(&chunk);
+                aes.decrypt_block(&mut block);
+                chunk.copy_from_slice(&block);
+            }
+
+            // Drop the padding
+            text.truncate(length);
+
+            Ok(Bytes::from(text))
+        } else {
+            Err(MeshcoreCryptoError::KeyLengthError)
+        }
+    }
 
-        // Bytes::copy_from_slice(&block)
-        ciphertext
+    pub fn encrypt(&self, plaintext: Bytes) -> Result<Bytes, MeshcoreCryptoError> {
+        use aes::cipher::KeyInit;
+
+        if let Ok(aes) = Aes256::new_from_slice(&self.0) {
+            let mut text = BytesMut::from(plaintext);
+
+            // The decryption function works on a 16-byte block of data.
+            // we need to break the input ciphertext into blocks of this size
+            // and work with them individually.  Also, any "short" blocks have
+            // to be zero-padded.
+            // Copy the original length so we can truncate back down
+            let chunk_size = 16;
+            let length = text.len();
+            let remainder = length % chunk_size;
+            // Pad the bytes to it's an even multiple of the chunk size
+
+            if remainder != 0 {
+                let padding = chunk_size - remainder;
+                text.reserve(padding);
+                text.put_bytes(0, padding);
+            }
+            
+            let chunks = text.chunks_exact_mut(chunk_size);
+            for chunk in chunks {
+                #[allow(deprecated)]
+                let mut block = *GenericArray::from_slice(&chunk);
+                aes.encrypt_block(&mut block);
+                chunk.copy_from_slice(&block);
+            }
+
+            Ok(Bytes::from(text))
+        } else {
+            Err(MeshcoreCryptoError::KeyLengthError)
+        }
     }
 }
 
@@ -118,17 +191,28 @@ impl FromStr for PublicKey {
 }
 
 impl FromStr for PrivateKey {
-    type Err = hex::FromHexError;
+    type Err = MeshcoreCryptoError;
 
     fn from_str(hex_str: &str) -> Result<Self, Self::Err> {
-        let hex = decode(hex_str)?;
-
-        if let Ok(bytes) = TryInto::<[u8; 64]>::try_into(hex) {
-            let exp_sec_key = ExpandedSecretKey::from_bytes(&bytes);
-
-            Ok(PrivateKey(exp_sec_key))
+        if let Ok(hex) = decode(hex_str) {
+            match hex.len() {
+                32 => {
+                    let bytes: [u8; 32] = hex.try_into().unwrap();
+                    let exp_sec_key = ExpandedSecretKey {
+                        scalar: Scalar::from_bytes_mod_order(bytes),
+                        hash_prefix: [0_u8; 32]
+                    };
+
+                    Ok(PrivateKey(exp_sec_key))
+                },
+                64 => {
+                    let bytes: [u8; 64] = hex.try_into().unwrap();
+                    Ok(PrivateKey(ExpandedSecretKey::from_bytes(&bytes)))
+                }
+                _ => Err(MeshcoreCryptoError::KeyLengthError)
+            }
         } else {
-            Err(hex::FromHexError::InvalidStringLength)
+            Err(MeshcoreCryptoError::HexDecodeError)
         }
     }
 }
@@ -165,9 +249,8 @@ pub struct Keychain {
 
 #[cfg(test)]
 mod tests {
-    // use curve25519_dalek::Scalar;
     use hex::{decode_to_slice, encode};
-
+    use rand_core::TryRngCore;
     use super::*;
 
     #[test]
@@ -193,22 +276,27 @@ mod tests {
     #[test]
     fn shared_secret() {
         // We'll make a public/private pair for alice and bob and make sure the shared secret is the same and expected
-        let alice_private = PrivateKey::from_str("4885CF25975EA09742EF76DA587D0957E74EE02AAA34A001458E207E63CF7E6C4940C8C42C335862C71CC2F139633057D1FEE5687B172B27E1E0302A1D480E08").unwrap();
-        let bob_private = PrivateKey::from_str("08976A389FA16B077492BA7403A2178F4DF22B74A44DA0BE780CD0A51F5796437BB76B51320EE216F483F741FD73ED32F7DF5BBCBE811F405E579DD45AA8280A").unwrap();
+        let alice_private = PrivateKey::from_str("58f5052c13275c8a3f4863a082555fed7ea08b9dec2eb00dd86f6b5412174458").unwrap();
+        let bob_private = PrivateKey::from_str("586c4cc29635af5865abe4f231bafbd373969725493c07271e02c7fec8ff3b5f").unwrap();
 
         let alice_public = PublicKey::from(&alice_private);
         let bob_public = PublicKey::from(&bob_private);
 
+        assert!(alice_public.0.as_bytes().to_vec() == decode("f673f022cf0466a95db086adf19084d4cf06e56fdabb202dc833bab227a78561").unwrap());
+        assert!(  bob_public.0.as_bytes().to_vec() == decode("47cb9867a8df6d2637eabe399b70f58f30f250da3911a4add950e7fac6aa3e39").unwrap());
+
         println!("Alice's public key: {}", encode(&alice_public.0.to_bytes()));
         println!("Bob's   public key: {}", encode(&bob_public.0.to_bytes()));
 
         let left_secret  = alice_private.create_secret(&bob_public);
         let right_secret = bob_private.create_secret(&alice_public);
 
+        assert!(left_secret.0.to_vec() == decode("ddbb8b5e70099817db83b48caa73f44a120b12a26072e5c29a023f16a2cd8b2a").unwrap());
+
         println!("Left shared secret:  {}", encode(left_secret.0));
         println!("Right shared secret: {}", encode(right_secret.0));
 
-        // assert!(left_secret == right_secret);
+        assert!(left_secret == right_secret);
     }
 
     #[test]
@@ -223,4 +311,127 @@ mod tests {
         let mac = group_secret.get_hmac(sample_data).unwrap();
         assert!(0xC3C1 == mac);
     }
+
+    #[test]
+    fn decrypt() {
+        let ciphertext = Bytes::copy_from_slice(&decode("354D619BAE9590E4D177DB7EEAF982F5BDCF78005D75157D9535FA90178F785D").unwrap());
+        let secret = SharedSecret::new_from_group_secret(Bytes::copy_from_slice(&decode("8b3387e9c5cdea6ac9e5edbaa115cd72").unwrap()));
+        let cleartext = secret.decrypt(ciphertext).unwrap();
+        println!("Cleartext: {}", encode(&cleartext));
+    }
+
+    #[test]
+    fn decrypt_online_example() {
+        let ciphertext = Bytes::copy_from_slice(&decode("95CCD78CB3FE2DE57774D552558AE954").unwrap());
+        let secret = SharedSecret::new_from_group_secret(Bytes::copy_from_slice(&decode("FB05C2FF72C9E7F931E5FA1232AFA962EC04367D016F10493F82884823A6B529").unwrap()));
+        let cleartext = secret.decrypt(ciphertext).unwrap();
+        let vec = cleartext.to_vec();
+        let string = String::from_utf8_lossy(&vec);
+        assert!("Hello my world!!" == string);
+    }
+
+    #[test]
+    fn encrypt_online_example() {
+        let plaintext = Bytes::copy_from_slice("Meshcore!".as_bytes());
+        let secret = SharedSecret::new_from_group_secret(Bytes::copy_from_slice(&decode("C8ED5DFD3DC316A7FA22450D6DC7097F52255F91E4E30F90F8CA3D25AC3E90E4").unwrap()));
+        let ciphertext = secret.encrypt(plaintext).unwrap();
+        assert!(ciphertext == decode("FE952E7104F0AC1CFC42CA5EE3783961").unwrap());
+    }
+
+    #[test]
+    fn chunks() {
+        // Test my understanding of how Byte's chunk_exact methods work for use in (en|de)cryption
+        let array = decode("0102030405060708090A0B0C0D0E0F").unwrap();
+        let mut text = BytesMut::from(array.as_slice());
+
+        // Copy the original length so we can truncate back down
+        let chunk_size = 4;
+        let length = text.len();
+
+        // Pad the bytes to it's an even multiple of the chunk size
+        let padding = chunk_size - (length % chunk_size);
+        text.reserve(padding);
+        text.put_bytes(0, padding);
+        
+        let chunks = text.chunks_exact_mut(4);
+        {
+            for chunk in chunks {
+                for item in chunk { *item = *item << 4; }
+            }
+        }
+
+        // Drop the padding
+        text.truncate(length);
+
+        println!("Result: {:#?}", encode(&text));
+
+        assert!(text == decode("102030405060708090A0B0C0D0E0F0").unwrap());
+        
+    }
+
+    // Example from the crate we're using
+    #[test]
+    fn aes_test() {
+        use aes::Aes128;
+        use aes::cipher::{
+            BlockCipher, BlockEncrypt, BlockDecrypt, KeyInit,
+            generic_array::GenericArray,
+        };
+
+        let key = GenericArray::from([0u8; 16]);
+        let mut block = GenericArray::from([42u8; 16]);
+
+        // Initialize cipher
+        let cipher = Aes128::new(&key);
+
+        let block_copy = block.clone();
+
+        // Encrypt block in-place
+        println!("Before: {}", encode(&block));
+        cipher.encrypt_block(&mut block);
+        println!("Crypted: {}", encode(&block));
+        // And decrypt it back
+        cipher.decrypt_block(&mut block);
+        println!("Decrypted: {}", encode(&block));
+
+        assert_eq!(block, block_copy);
+
+        // Implementation supports parallel block processing. Number of blocks
+        // processed in parallel depends in general on hardware capabilities.
+        // This is achieved by instruction-level parallelism (ILP) on a single
+        // CPU core, which is differen from multi-threaded parallelism.
+        let mut blocks = [block; 100];
+        cipher.encrypt_blocks(&mut blocks);
+
+        for block in blocks.iter_mut() {
+            cipher.decrypt_block(block);
+            assert_eq!(block, &block_copy);
+        }
+
+        // `decrypt_blocks` also supports parallel block processing.
+        cipher.decrypt_blocks(&mut blocks);
+
+        for block in blocks.iter_mut() {
+            cipher.encrypt_block(block);
+            assert_eq!(block, &block_copy);
+        }
+    }
+
+    #[test]
+    fn shared_secrets_example() {
+        use x25519_dalek::{StaticSecret, PublicKey};
+
+        let a_s: [u8; 32] = decode("58f5052c13275c8a3f4863a082555fed7ea08b9dec2eb00dd86f6b5412174458").unwrap().try_into().unwrap();
+        let b_s: [u8; 32] = decode("586c4cc29635af5865abe4f231bafbd373969725493c07271e02c7fec8ff3b5f").unwrap().try_into().unwrap();
+        let alice_secret = StaticSecret::from(a_s);
+        let alice_public = PublicKey::from(&alice_secret);
+
+        let bob_secret = StaticSecret::from(b_s);
+        let bob_public = PublicKey::from(&bob_secret);
+
+        let alice_shared_secret = alice_secret.diffie_hellman(&bob_public);
+        let bob_shared_secret = bob_secret.diffie_hellman(&alice_public);
+        
+        assert_eq!(alice_shared_secret.as_bytes(), bob_shared_secret.as_bytes());
+    }
 }
\ No newline at end of file
diff --git a/src/packet.rs b/src/packet.rs
index 5788100..cd403fd 100644
--- a/src/packet.rs
+++ b/src/packet.rs
@@ -173,7 +173,7 @@ enum PayloadVersion {
 impl From<u8> for PayloadVersion {
     fn from(value: u8) -> Self {
         let value = (value & 0xC0) >> 6;
-        assert!(value < 4, "Programming error in masking and bitshifting");
+        assert!(value < 4, "Programming error in masking and bit-shifting");
 
         match value {
             0x00 => PayloadVersion::VersionOne,
@@ -407,7 +407,7 @@ impl From<Bytes> for Advert {
 
         // Lastly the name...  The flag is a little
         // irrelevant because the only different is
-        // whether the name is ommitted or just empty.
+        // whether the name is omitted or just empty.
         // I'm not going to make the distinction
 
         // The string is assumed to be utf8, but if that