diff --git a/src/netmd/encryption.rs b/src/netmd/encryption.rs index fb1b9ab..944870c 100644 --- a/src/netmd/encryption.rs +++ b/src/netmd/encryption.rs @@ -9,11 +9,10 @@ use super::interface::DataEncryptorInput; type DesEcbEnc = ecb::Decryptor; type DesCbcEnc = cbc::Encryptor; -pub fn new_thread_encryptor( - _input: DataEncryptorInput, +pub fn threaded_encryptor( + input: DataEncryptorInput, ) -> UnboundedReceiver<(Vec, Vec, Vec)> { let (tx, rx) = unbounded_channel::<(Vec, Vec, Vec)>(); - let input = Box::from(_input); thread::spawn(move || { let mut iv = [0u8; 8]; @@ -76,3 +75,103 @@ pub fn new_thread_encryptor( rx } + +pub struct Encryptor { + input_data: Vec, + iv: [u8; 8], + random_key: [u8; 8], + encrypted_random_key: [u8; 8], + default_chunk_size: usize, + current_chunk_size: usize, + offset: usize, + packet_count: usize, + closed: bool, +} + +impl Encryptor { + pub fn new(input: DataEncryptorInput) -> Self { + let iv = [0u8; 8]; + + // Create the random key + let mut random_key = [0u8; 8]; + rand::thread_rng().fill_bytes(&mut random_key); + + // Encrypt it with the kek + let mut encrypted_random_key = random_key; + if let Err(x) = DesEcbEnc::new(&input.kek.into()) + .decrypt_padded_mut::(&mut encrypted_random_key) + { + panic!("Cannot create main key {:?}", x) + }; + + let default_chunk_size = match input.chunk_size { + 0 => 0x00100000, + e => e, + }; + + let packet_count = 0; + let current_chunk_size = 0; + + let mut input_data = input.data.clone(); + if (input_data.len() % input.frame_size) != 0 { + let padding_remaining = input.frame_size - (input_data.len() % input.frame_size); + input_data.extend(std::iter::repeat(0).take(padding_remaining)); + } + + let offset: usize = 0; + + Self { + input_data, + iv, + random_key, + encrypted_random_key, + current_chunk_size, + offset, + default_chunk_size, + packet_count, + closed: false, + } + } + + /// Get the next encrypted value + pub async fn next(&mut self) -> Option<(Vec, Vec, Vec)> { + if self.closed { + return None + } + + if self.packet_count > 0 { + self.current_chunk_size = self.default_chunk_size; + } else { + self.current_chunk_size = self.default_chunk_size - 24; + } + + self.current_chunk_size = std::cmp::min(self.current_chunk_size, self.input_data.len() - self.offset); + + let this_data_chunk = &mut self.input_data[self.offset..self.offset + self.current_chunk_size]; + DesCbcEnc::new(&self.random_key.into(), &self.iv.into()) + .encrypt_padded_mut::(this_data_chunk, self.current_chunk_size) + .unwrap(); + + let output = ( + self.encrypted_random_key.to_vec(), + self.iv.to_vec(), + this_data_chunk.to_vec(), + ); + + self.iv.copy_from_slice(&this_data_chunk[this_data_chunk.len() - 8..]); + + self.packet_count += 1; + self.offset += self.current_chunk_size; + + Some(output) + } + + /// Call close to return none from subsequent calls + pub fn close(&mut self) { + self.closed = true; + } +} + +pub fn encryptor() { + +} diff --git a/src/netmd/interface.rs b/src/netmd/interface.rs index 6057276..7be62ec 100644 --- a/src/netmd/interface.rs +++ b/src/netmd/interface.rs @@ -15,9 +15,9 @@ use std::collections::HashMap; use std::error::Error; use std::time::Duration; use thiserror::Error; -use tokio::sync::mpsc::UnboundedReceiver; use super::base::NetMD; +use super::encryption::{threaded_encryptor, EncryptorState}; use super::utils::{cross_sleep, to_sjis}; /// An action to take on the player @@ -1691,7 +1691,10 @@ impl NetMDInterface { frames: u32, pkt_size: u32, // key, iv, data + #[cfg(not(target_family = "wasm"))] mut packets: UnboundedReceiver<(Vec, Vec, Vec)>, + #[cfg(target_family = "wasm")] + mut packets: EncryptorState, hex_session_key: &[u8], progress_callback: F, ) -> Result<(u16, Vec, Vec), InterfaceError> @@ -1863,10 +1866,6 @@ pub struct MDTrack { pub data: Vec, pub chunk_size: usize, pub full_width_title: Option, - - #[allow(clippy::type_complexity)] - pub encrypt_packets_iterator: - Box UnboundedReceiver<(Vec, Vec, Vec)>>, } pub struct DataEncryptorInput { @@ -1921,8 +1920,19 @@ impl MDTrack { [0x14, 0xe3, 0x83, 0x4e, 0xe2, 0xd3, 0xcc, 0xa5] } + #[cfg(not(target_family = "wasm"))] pub fn get_encrypting_iterator(&mut self) -> UnboundedReceiver<(Vec, Vec, Vec)> { - (self.encrypt_packets_iterator)(DataEncryptorInput { + threaded_encryptor(DataEncryptorInput { + kek: self.get_kek(), + frame_size: self.frame_size(), + chunk_size: self.chunk_size(), + data: std::mem::take(&mut self.data), + }) + } + + #[cfg(target_family = "wasm")] + pub fn get_encrypting_iterator(&mut self) -> EncryptorState { + EncryptorState::new(DataEncryptorInput { kek: self.get_kek(), frame_size: self.frame_size(), chunk_size: self.chunk_size(),