diff --git a/src/encrypt.rs b/src/encrypt.rs index f9311ac..eff5b90 100644 --- a/src/encrypt.rs +++ b/src/encrypt.rs @@ -2,10 +2,90 @@ use rsa::PublicKey; use rsa::{RsaPrivateKey, RsaPublicKey, PaddingScheme, errors::Result}; -use rand::rngs::OsRng; -use aes::Aes128; +use rand::{Rng, rngs::OsRng}; +use aes::{Aes128, NewBlockCipher}; use aes::cipher::{BlockEncrypt, BlockDecrypt, generic_array::GenericArray}; +#[derive(Clone)] +pub struct McCipher { + key: [u8; 16], + state_en: [u8; 16], + state_de: [u8; 16], +} + +impl McCipher { + pub fn create() -> Self { + let mut rng = rand::thread_rng(); + let aes_key: [u8; 16] = rng.gen(); + Self { + key: aes_key.clone(), + state_en: aes_key.clone(), + state_de: aes_key.clone(), + } + } + + pub fn get_encrypted_key( + &self, + public_key: &RsaPublicKey, + ) -> Result> { + encrypt_rsa(public_key, &self.key) + } + + pub fn create_with_encrypted_key( + private_key: &RsaPrivateKey, + data: &[u8], + ) -> Result { + let aes_key: [u8; 16] = decrypt_rsa(private_key, data)? + .as_slice()[0..16].try_into().unwrap(); + Ok(Self { + key: aes_key.clone(), + state_en: aes_key.clone(), + state_de: aes_key.clone(), + }) + } + + pub fn encrypt_aes(&self, data: Vec) -> Vec { + let mut out_data = vec![0; data.len()]; + for i in 0..data.len() { + out_data[i] = self.encrypt_block(data[i]); + } + out_data + } + + pub fn decrypt_aes(&self, data: Vec) -> Vec { + let mut out_data = vec![0; data.len()]; + for i in 0..data.len() { + out_data[i] = self.decrypt_block(data[i]); + } + out_data + } + + fn shift_left(mut arr: [u8; 16], new: u8) { + for i in 0..arr.len() - 1 { + arr[i] = arr[i + 1]; + } + arr[arr.len() - 1] = new; + } + + fn encrypt_block(&self, data: u8) -> u8 { + let cipher = Aes128::new(GenericArray::from_slice(&self.key)); + let mut block = GenericArray::clone_from_slice(&self.state_en); + cipher.encrypt_block(&mut block); + let data = data ^ block[0]; + Self::shift_left(self.state_en, data); + data + } + + fn decrypt_block(&self, data: u8) -> u8 { + let cipher = Aes128::new(GenericArray::from_slice(&self.key)); + let mut block = GenericArray::clone_from_slice(&self.state_en); + cipher.decrypt_block(&mut block); + Self::shift_left(self.state_de, data); + let data = data ^ block[0]; + data + } +} + pub fn generate_rsa_keys() -> Result { let mut rng = OsRng; let bits = 2048; @@ -29,15 +109,3 @@ pub fn decrypt_rsa( let padding = PaddingScheme::new_pkcs1v15_encrypt(); private_key.decrypt(padding, data) } - -pub fn encrypt_aes(cipher: Aes128, data: &[u8; 16]) -> Vec { - let mut block = GenericArray::clone_from_slice(data); - cipher.encrypt_block(&mut block); - block.to_vec() -} - -pub fn decrypt_aes(cipher: Aes128, data: &[u8; 16]) -> Vec { - let mut block = GenericArray::clone_from_slice(data); - cipher.decrypt_block(&mut block); - block.to_vec() -} diff --git a/src/mc_types.rs b/src/mc_types.rs index 3643cb1..0f0640f 100644 --- a/src/mc_types.rs +++ b/src/mc_types.rs @@ -10,11 +10,9 @@ use async_trait::async_trait; use rsa::{RsaPrivateKey, RsaPublicKey}; use rsa::pkcs8::{EncodePublicKey, DecodePublicKey}; use rand::Rng; -use aes::Aes128; -use aes::cipher::{NewBlockCipher, generic_array::GenericArray}; use crate::login; -use crate::encrypt; +use crate::encrypt::{self, McCipher}; use crate::play::Play; pub type Result = std::result::Result>; @@ -70,8 +68,7 @@ pub struct ProtocolConnection<'a> { pub stream_write: &'a mut OwnedWriteHalf, rsa_private_key: Option, rsa_public_key: Option, - aes_encryption_key: Option<[u8; 16]>, - aes_cipher: Option, + aes_cipher: Option, verify_token: Option<[u8; 16]>, } @@ -85,7 +82,6 @@ impl<'a> ProtocolConnection<'a> { stream_write, rsa_private_key: None, rsa_public_key: None, - aes_encryption_key: None, aes_cipher: None, verify_token: None, } @@ -130,17 +126,14 @@ impl<'a> ProtocolConnection<'a> { ) -> Result { self.rsa_public_key = Some( RsaPublicKey::from_public_key_der(&request.public_key)?); - let mut rng = rand::thread_rng(); - let aes_key = rng.gen(); - self.aes_encryption_key = Some(aes_key); - self.aes_cipher = Some(Aes128::new(GenericArray::from_slice(&aes_key))); - match self.aes_encryption_key { - Some(key) => { + self.aes_cipher = Some(McCipher::create()); + match &self.aes_cipher { + Some(aes_cipher) => { match &self.rsa_public_key { Some(public_key) => { Ok(login::serverbound::EncryptionResponse { - shared_secret: encrypt::encrypt_rsa( - public_key, &key)?, + shared_secret: aes_cipher + .get_encrypted_key(public_key)?, verify_token: encrypt::encrypt_rsa( public_key, request.verify_token[0..16] @@ -168,11 +161,11 @@ impl<'a> ProtocolConnection<'a> { &private_key, response.verify_token.as_slice() )? == token { - self.aes_encryption_key = - Some(encrypt::decrypt_rsa( - &private_key, - response.shared_secret.as_slice() - )?[0..16].try_into().unwrap()); + self.aes_cipher = + Some(McCipher::create_with_encrypted_key( + private_key, + response.shared_secret.as_slice(), + )?); Ok(()) } else { Err(Box::new(PacketError::EncryptionError)) @@ -216,27 +209,14 @@ unsafe impl<'a> Send for ProtocolConnection<'a> {} #[async_trait] impl<'a> ProtocolRead for ProtocolConnection<'a> { async fn read_data(&mut self) -> Result> { - match self.aes_cipher.clone() { + match &self.aes_cipher { Some(aes_cipher) => { - let mut buffer: Vec = vec![0; 16]; + let length = read_var_int_stream_encrypted( + self.stream_read, aes_cipher).await? as usize; + + let mut buffer: Vec = vec![0; length]; self.stream_read.read_exact(&mut buffer).await?; - buffer = encrypt::decrypt_aes( - aes_cipher, buffer[0..16].try_into().unwrap()); - let raw_length = read_var_int_vec(&mut buffer)?; - let length = - if (raw_length - buffer.len() as i32) % 16 == 0 { - (raw_length - buffer.len() as i32) / 16 - } else { - ((raw_length - buffer.len() as i32) / 16) + 1 - }; - - for _ in 0..length { - let mut block: Vec = vec![0; 16]; - self.stream_read.read_exact(&mut block).await?; - buffer.append(&mut block); - } - - Ok(buffer) + Ok(aes_cipher.decrypt_aes(buffer)) }, None => { let length = read_var_int_stream( @@ -256,21 +236,10 @@ impl<'a> ProtocolWrite for ProtocolConnection<'a> { async fn write_data(&mut self, data: &mut Vec) -> Result<()> { let mut out_data = convert_var_int(data.len() as i32); out_data.append(data); - match self.aes_cipher.clone() { + match &self.aes_cipher { Some(aes_cipher) => { - let length = - if (data.len() as i32) % 16 == 0 { - (data.len() as i32) / 16 - } else { - ((data.len() as i32) / 16) + 1 - }; - - for _ in 0..length { - let mut block: Vec = out_data[0..16].to_vec(); - block = encrypt::encrypt_aes( - aes_cipher.clone(), block[0..16].try_into().unwrap()); - self.stream_write.write_all(&block).await?; - } + self.stream_write.write_all( + &aes_cipher.encrypt_aes(out_data)).await?; Ok(()) }, @@ -285,7 +254,7 @@ impl<'a> ProtocolWrite for ProtocolConnection<'a> { pub struct WriteHaftProtocolConnection<'a> { pub stream_write: &'a mut OwnedWriteHalf, - aes_cipher: Option, + aes_cipher: Option, } impl<'a> WriteHaftProtocolConnection<'a> { @@ -303,28 +272,13 @@ unsafe impl<'a> Send for WriteHaftProtocolConnection<'a> {} #[async_trait] impl<'a> ProtocolWrite for WriteHaftProtocolConnection<'a> { - async fn write_data( - &mut self, - data: &mut Vec, - ) -> Result<()> { + async fn write_data(&mut self, data: &mut Vec) -> Result<()> { let mut out_data = convert_var_int(data.len() as i32); out_data.append(data); - match self.aes_cipher.clone() { + match &self.aes_cipher { Some(aes_cipher) => { - let length = - if (data.len() as i32) % 16 == 0 { - (data.len() as i32) / 16 - } else { - ((data.len() as i32) / 16) + 1 - }; - - for _ in 0..length { - let mut block: Vec = out_data[0..16].to_vec(); - block = encrypt::encrypt_aes( - aes_cipher.clone(), block[0..16].try_into().unwrap()); - self.stream_write.write_all(&block).await?; - } - + self.stream_write.write_all( + &aes_cipher.encrypt_aes(out_data)).await?; Ok(()) }, @@ -339,7 +293,7 @@ impl<'a> ProtocolWrite for WriteHaftProtocolConnection<'a> { pub struct ReadHaftProtocolConnection<'a> { pub stream_read: &'a mut OwnedReadHalf, - aes_cipher: Option, + aes_cipher: Option, } impl<'a> ReadHaftProtocolConnection<'a> { @@ -370,27 +324,14 @@ unsafe impl<'a> Send for ReadHaftProtocolConnection<'a> {} #[async_trait] impl<'a> ProtocolRead for ReadHaftProtocolConnection<'a> { async fn read_data(&mut self) -> Result> { - match self.aes_cipher.clone() { + match &self.aes_cipher { Some(aes_cipher) => { - let mut buffer: Vec = vec![0; 16]; + let length = read_var_int_stream_encrypted( + self.stream_read, aes_cipher).await? as usize; + + let mut buffer: Vec = vec![0; length]; self.stream_read.read_exact(&mut buffer).await?; - buffer = encrypt::decrypt_aes( - aes_cipher, buffer[0..16].try_into().unwrap()); - let raw_length = read_var_int_vec(&mut buffer)?; - let length = - if (raw_length - buffer.len() as i32) % 16 == 0 { - (raw_length - buffer.len() as i32) / 16 - } else { - ((raw_length - buffer.len() as i32) / 16) + 1 - }; - - for _ in 0..length { - let mut block: Vec = vec![0; 16]; - self.stream_read.read_exact(&mut block).await?; - buffer.append(&mut block); - } - - Ok(buffer) + Ok(aes_cipher.decrypt_aes(buffer)) }, None => { let length = read_var_int_stream( @@ -443,6 +384,27 @@ async fn read_var_int_stream(stream: &mut OwnedReadHalf) -> Result { Ok(varint) } +async fn read_var_int_stream_encrypted( + stream: &mut OwnedReadHalf, + cipher: &McCipher, +) -> Result { + let mut data: Vec = vec![]; + + loop { + let encrypted_byte = stream.read_u8().await?; + let mut current_byte = cipher.decrypt_aes(vec![encrypted_byte]); + + data.append(&mut current_byte); + + if (current_byte[0] & CONTINUE_BIT) == 0 { + break; + } + } + + let varint = get_var_int(&mut data)?; + + Ok(varint) +} fn read_var_int_vec(stream: &mut Vec) -> Result { let mut data: Vec = vec![];