From ffaae0e3ae705d78e570ddc71881c2bfded8b897 Mon Sep 17 00:00:00 2001
From: arkpar <arkady.paronyan@gmail.com>
Date: Wed, 17 Aug 2016 21:02:59 +0200
Subject: [PATCH] ECIES encrypt/decrypt for a single message

---
 util/src/crypto.rs | 69 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/util/src/crypto.rs b/util/src/crypto.rs
index 8930ed42a59..e96ea3a7c41 100644
--- a/util/src/crypto.rs
+++ b/util/src/crypto.rs
@@ -351,6 +351,34 @@ pub mod ecies {
 		Ok(msg)
 	}
 
+	/// Encrypt a message with a public key
+	pub fn encrypt_single_message(public: &Public, plain: &[u8]) -> Result<Bytes, CryptoError> {
+		use ::rcrypto::digest::Digest;
+		use ::rcrypto::sha2::Sha256;
+		let r = try!(KeyPair::create());
+		let z = try!(ecdh::agree(r.secret(), public));
+		let mut key = [0u8; 32];
+		let mut mkey = [0u8; 32];
+		kdf(&z, &[0u8; 0], &mut key);
+		let mut hasher = Sha256::new();
+		let mkey_material = &key[16..32];
+		hasher.input(mkey_material);
+		hasher.result(&mut mkey);
+		let ekey = &key[0..16];
+
+		let mut msgd = vec![0u8; (64 + 16 + plain.len())];
+		{
+			r.public().copy_to(&mut msgd[0..64]);
+			let iv = H128::random();
+			iv.copy_to(&mut msgd[64..(64+16)]);
+			{
+				let cipher = &mut msgd[(64 + 16)..(64 + 16 + plain.len())];
+				aes::encrypt(ekey, &iv, plain, cipher);
+			}
+		}
+		Ok(msgd)
+	}
+
 	/// Decrypt a message with a secret key
 	pub fn decrypt(secret: &Secret, shared_mac: &[u8], encrypted: &[u8]) -> Result<Bytes, CryptoError> {
 		use ::rcrypto::digest::Digest;
@@ -396,6 +424,37 @@ pub mod ecies {
 		Ok(msg)
 	}
 
+	/// Decrypt single message with a secret key
+	pub fn decrypt_single_message(secret: &Secret, encrypted: &[u8]) -> Result<Bytes, CryptoError> {
+		use ::rcrypto::digest::Digest;
+		use ::rcrypto::sha2::Sha256;
+
+		let meta_len = 64 + 16;
+		if encrypted.len() < meta_len {
+			return Err(CryptoError::InvalidMessage); //invalid message: publickey
+		}
+
+		let e = encrypted;
+		let p = Public::from_slice(&e[0..64]);
+		let z = try!(ecdh::agree(secret, &p));
+		let mut key = [0u8; 32];
+		kdf(&z, &[0u8; 0], &mut key);
+		let ekey = &key[0..16];
+		let mkey_material = &key[16..32];
+		let mut hasher = Sha256::new();
+		let mut mkey = [0u8; 32];
+		hasher.input(mkey_material);
+		hasher.result(&mut mkey);
+
+		let clen = encrypted.len() - meta_len;
+		let cipher_with_iv = &e[64..(64+16+clen)];
+		let cipher_iv = &cipher_with_iv[0..16];
+		let cipher_no_iv = &cipher_with_iv[16..];
+		let mut msg = vec![0u8; clen];
+		aes::decrypt(ekey, cipher_iv, cipher_no_iv, &mut msg[..]);
+		Ok(msg)
+	}
+
 	fn kdf(secret: &Secret, s1: &[u8], dest: &mut [u8]) {
 		use ::rcrypto::digest::Digest;
 		use ::rcrypto::sha2::Sha256;
@@ -489,4 +548,14 @@ mod tests {
 		let decrypted = ecies::decrypt(kp.secret(), shared, &encrypted).unwrap();
 		assert_eq!(decrypted[..message.len()], message[..]);
 	}
+
+	#[test]
+	fn ecies_shared_single() {
+		let kp = KeyPair::create().unwrap();
+		let message = b"So many books, so little time";
+		let encrypted = ecies::encrypt_single_message(kp.public(), message).unwrap();
+		assert!(encrypted[..] != message[..]);
+		let decrypted = ecies::decrypt_single_message(kp.secret(), &encrypted).unwrap();
+		assert_eq!(decrypted[..message.len()], message[..]);
+	}
 }