add epk3 initial support , move epk common functions

src/formats/epk.rs

@@ -51,8 +51,56 @@ pub fn extract_epk(file: &File, output_folder: &str) -> Result<(), Box<dyn std::
         println!("EPK2 detected!\n");
         formats::epk2::extract_epk2(file, output_folder)?;
     } else if epk_version == Some("epk3".to_string()) {
-        println!("EPK3 detected! Not supported yet.");
+        println!("EPK3 detected!\n");
+        formats::epk3::extract_epk3(file, output_folder)?;
     }
 
     Ok(())
+}
+
+//COMMON EPK FUNCTIONS
+pub fn find_key<'a>(key_array: &'a [(&'a str, &'a str)], data: &[u8], expected_magic: &[u8]) -> Result<Option<(&'a str, Vec<u8>)>, Box<dyn std::error::Error>> {
+    for (key_hex, name) in key_array {
+        let key_bytes = hex::decode(key_hex)?;
+        let decrypted = match decrypt_aes_ecb_auto(&key_bytes, data) {
+            Ok(d) => d,
+            Err(_) => continue,
+        };
+     
+        if decrypted.starts_with(expected_magic) {
+            return Ok(Some((name, key_bytes)));
+        }
+    }
+    Ok(None)
+}
+
+use aes::Aes128;
+use aes::Aes256;
+use ecb::{Decryptor, cipher::{BlockDecryptMut, KeyInit, generic_array::GenericArray}};
+
+type Aes128EcbDec = Decryptor<Aes128>;
+type Aes256EcbDec = Decryptor<Aes256>;
+
+pub fn decrypt_aes_ecb_auto(key: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
+    let mut buffer = ciphertext.to_vec();
+
+    if key.len() == 32 {
+        // aes256
+        let key_array: [u8; 32] = key.try_into()?;
+        let mut decryptor = Aes256EcbDec::new(&key_array.into());
+        for chunk in buffer.chunks_exact_mut(16) {
+            let block: &mut [u8; 16] = chunk.try_into()?;
+            decryptor.decrypt_block_mut(GenericArray::from_mut_slice(block));
+        }
+    } else {
+        // aes128
+        let key_array: [u8; 16] = key.try_into()?;
+        let mut decryptor = Aes128EcbDec::new(&key_array.into());
+        for chunk in buffer.chunks_exact_mut(16) {
+            let block: &mut [u8; 16] = chunk.try_into()?;
+            decryptor.decrypt_block_mut(GenericArray::from_mut_slice(block));
+        }
+    }
+
+    Ok(buffer)
 }

src/formats/epk2.rs

@@ -2,11 +2,9 @@ use std::fs::{self, File, OpenOptions};
 use std::path::{Path};
 use std::io::{Write, Seek, SeekFrom, Cursor};
 
-use aes::Aes128;
-use ecb::{Decryptor, cipher::{BlockDecryptMut, KeyInit, generic_array::GenericArray}};
-
 use crate::common;
 use crate::keys;
+use crate::formats::epk::{decrypt_aes_ecb_auto, find_key};
 
 pub fn is_epk2_file(file: &File) -> bool {
     let header = common::read_file(&file, 128, 4).expect("Failed to read from file.");
@@ -23,35 +21,6 @@ struct Pak {
     name: String,
 }
 
-fn find_key<'a>(data: &[u8], expected_magic: &[u8]) -> Result<Option<(&'a str, Vec<u8>)>, Box<dyn std::error::Error>> {
-    for (key_hex, name) in keys::EPK2 {
-        let key_bytes = hex::decode(key_hex)?;
-        let decrypted = match decrypt_aes128_ecb(&key_bytes, data) {
-            Ok(d) => d,
-            Err(_) => continue,
-        };
-        if decrypted.starts_with(expected_magic) {
-            return Ok(Some((name, key_bytes)));
-        }
-    }
-    Ok(None)
-}
-
-type Aes128EcbDec = Decryptor<Aes128>;
-
-fn decrypt_aes128_ecb(key: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
-    let key_array: [u8; 16] = key.try_into()?;
-    let mut decryptor = Aes128EcbDec::new(&key_array.into());
-    let mut buffer = ciphertext.to_vec();
-
-    for chunk in buffer.chunks_exact_mut(16) {
-        let block: &mut [u8; 16] = chunk.try_into()?;
-        decryptor.decrypt_block_mut(GenericArray::from_mut_slice(block));
-    }
-    
-    Ok(buffer)
-}
-
 pub fn extract_epk2(mut file: &File, output_folder: &str) -> Result<(), Box<dyn std::error::Error>> {
     file.seek(SeekFrom::Start(128))?; //inital signature
 
@@ -69,10 +38,10 @@ pub fn extract_epk2(mut file: &File, output_folder: &str) -> Result<(), Box<dyn
         println!("Header is encrypted...");
         println!("\nFinding key...");
         //find the key, knowing that the header should start with "epak"
-        if let Some((key_name, key_bytes)) = find_key(&stored_header, b"epak")? {
+        if let Some((key_name, key_bytes)) = find_key(&keys::EPK2, &stored_header, b"epak")? {
             println!("Found valid key: {}", key_name);
             matching_key = Some(key_bytes);
-            header = decrypt_aes128_ecb(matching_key.as_ref().unwrap(), &stored_header)?;
+            header = decrypt_aes_ecb_auto(matching_key.as_ref().unwrap(), &stored_header)?;
         } else {
             println!("No valid key found!");
             return Ok(());
@@ -137,7 +106,7 @@ pub fn extract_epk2(mut file: &File, output_folder: &str) -> Result<(), Box<dyn
         if matching_key.is_none() {
             println!("\nFinding key...");
             //find the key, knowing that the header should start with with the paks name
-            if let Some((key_name, key_bytes)) = find_key(&encrypted_header, pak.name.as_bytes())? {
+            if let Some((key_name, key_bytes)) = find_key(&keys::EPK2, &encrypted_header, pak.name.as_bytes())? {
                 println!("Found correct key: {}", key_name);
                 matching_key = Some(key_bytes);
             } else {
@@ -148,7 +117,7 @@ pub fn extract_epk2(mut file: &File, output_folder: &str) -> Result<(), Box<dyn
 
         let matching_key_bytes = matching_key.as_ref().unwrap();
 
-        let header = decrypt_aes128_ecb(&matching_key_bytes, &encrypted_header)?;
+        let header = decrypt_aes_ecb_auto(&matching_key_bytes, &encrypted_header)?;
 
         let segment_count = u32::from_le_bytes(header[84..88].try_into().unwrap());
         let mut segment_size = u32::from_le_bytes(header[88..92].try_into().unwrap());
@@ -162,7 +131,7 @@ pub fn extract_epk2(mut file: &File, output_folder: &str) -> Result<(), Box<dyn
                 signature_count += 1;
 
                 let encrypted_header = common::read_exact(&mut file, 128)?;
-                let header = decrypt_aes128_ecb(&matching_key_bytes, &encrypted_header)?;
+                let header = decrypt_aes_ecb_auto(&matching_key_bytes, &encrypted_header)?;
                 segment_size = u32::from_le_bytes(header[88..92].try_into().unwrap());
             }
 
@@ -186,7 +155,7 @@ pub fn extract_epk2(mut file: &File, output_folder: &str) -> Result<(), Box<dyn
             };
 
             let segment_data = common::read_exact(&mut file, actual_segment_size as usize)?;
-            let out_data = decrypt_aes128_ecb(&matching_key_bytes, &segment_data)?;
+            let out_data = decrypt_aes_ecb_auto(&matching_key_bytes, &segment_data)?;
 
             println!("- Segment {}/{}, size: {}", i + 1, segment_count, actual_segment_size);
 

src/formats/epk3.rs

@@ -0,0 +1,139 @@
+use std::fs::{self, File, OpenOptions};
+use std::path::{Path};
+use std::io::{Write, Seek, SeekFrom, Cursor};
+
+use crate::common;
+use crate::keys;
+use crate::formats::epk::{decrypt_aes_ecb_auto, find_key};
+
+pub fn extract_epk3(mut file: &File, output_folder: &str) -> Result<(), Box<dyn std::error::Error>> {
+    file.seek(SeekFrom::Start(128))?; //inital signature
+
+    let stored_header = common::read_exact(&mut file, 1584)?; //max header size
+    let header: Vec<u8>;
+
+    let mut new_type = false;
+
+    let mut matching_key: Option<Vec<u8>> = None;
+    println!("Finding key...");
+
+    // find the key, knowing that the header should start with "EPK3"
+    if let Some((key_name, key_bytes)) = find_key(&keys::EPK3, &stored_header, b"EPK3")? {
+        println!("Found valid key: {}", key_name);
+        matching_key = Some(key_bytes);
+        header = decrypt_aes_ecb_auto(matching_key.as_ref().unwrap(), &stored_header)?;
+
+    //try for new format epk3 where theres an additional 128byte signature at the beginning
+    } else if let Some((key_name, key_bytes)) = find_key(&keys::EPK3, &stored_header[128..], b"EPK3")? {
+        println!("Found valid key: {}", key_name);
+        matching_key = Some(key_bytes);
+        header = decrypt_aes_ecb_auto(matching_key.as_ref().unwrap(), &stored_header)?;
+        new_type = true;
+
+    } else {
+        println!("No valid key found!");
+        return Ok(());
+    }
+
+    let signature_size = if new_type {256} else {128};
+    let extra_segment_size = if new_type {4} else {0};
+
+    let matching_key_bytes = matching_key.as_ref().unwrap();
+
+    //parse header
+    let mut hdr_reader = Cursor::new(header);
+
+    if new_type {let _signature = common::read_exact(&mut hdr_reader, 128)?;};
+
+    let _epk3 = common::read_exact(&mut hdr_reader, 4)?; //EPK3 magic
+
+    let version = common::read_exact(&mut hdr_reader, 4)?;
+
+    let ota_id_bytes = common::read_exact(&mut hdr_reader, 32)?;
+    let ota_id = common::string_from_bytes(&ota_id_bytes);
+
+    let package_info_size_bytes = common::read_exact(&mut hdr_reader, 4)?;
+    let package_info_size = u32::from_le_bytes(package_info_size_bytes.try_into().unwrap());
+
+    println!("\nEPK info:\nOTA ID: {}\nVersion: {:02x?}.{:02x?}.{:02x?}\nPackage Info size: {}\n", 
+                ota_id, version[3], version[2], version[1], package_info_size);
+    //
+
+    let _versions = common::read_exact(&mut file, 36)?;
+
+    let _signature = common::read_exact(&mut file, signature_size)?;
+
+    //PKG INFO
+    let pkg_info_encrypted = common::read_exact(&mut file, package_info_size as usize)?;
+    let pkg_info = decrypt_aes_ecb_auto(matching_key_bytes, &pkg_info_encrypted)?;
+
+    let mut pkg_info_reader = Cursor::new(pkg_info);
+
+    let package_info_list_size_b = common::read_exact(&mut pkg_info_reader, 4)?;
+    let package_info_list_size = u32::from_le_bytes(package_info_list_size_b.try_into().unwrap());
+
+    let package_info_count_b = common::read_exact(&mut pkg_info_reader, 4)?;
+    let package_info_count = u32::from_le_bytes(package_info_count_b.try_into().unwrap());
+
+    println!("Package info list size: {}\nPackage info count: {}", 
+                package_info_list_size, package_info_count);
+
+    if new_type {let _unknown = common::read_exact(&mut pkg_info_reader, 4)?;}; //uncertain if this is only in new type, but i think it is
+
+    while (pkg_info_reader.position() as usize) < pkg_info_reader.get_ref().len() {
+        let segment = common::read_exact(&mut pkg_info_reader, 324)?;
+
+        let package_name_b = &segment[8..136];
+        let package_name = common::string_from_bytes(&package_name_b);
+
+        let package_size_b = &segment[296..300];
+        let package_size = u32::from_le_bytes(package_size_b.try_into().unwrap());
+
+        //Package segment info
+        let segment_index_b = &segment[308..312];
+        let mut segment_index = u32::from_le_bytes(segment_index_b.try_into().unwrap());
+
+        let segment_count_b = &segment[312..316];
+        let segment_count = u32::from_le_bytes(segment_count_b.try_into().unwrap());
+
+        let segment_size_b = &segment[316..320];
+        let mut segment_size = u32::from_le_bytes(segment_size_b.try_into().unwrap());
+        //
+
+        println!("\nPak - {}, Size: {}, Segment Count: {}",
+                package_name, package_size, segment_count);
+        
+        for i in 0..segment_count {
+            if i > 0 {
+                let segment = common::read_exact(&mut pkg_info_reader, 324)?;
+
+                let segment_index_b = &segment[308..312];
+                segment_index = u32::from_le_bytes(segment_index_b.try_into().unwrap());
+
+                let segment_size_b = &segment[316..320];
+                segment_size = u32::from_le_bytes(segment_size_b.try_into().unwrap()); 
+            }   
+            
+            println!("- Segment {}/{}, Size: {}", segment_index + 1, segment_count, segment_size);
+
+            let _signature = common::read_exact(&mut file, signature_size)?;
+
+            let encrypted_data = common::read_exact(&mut file, segment_size as usize + extra_segment_size)?;
+            let out_data = decrypt_aes_ecb_auto(matching_key_bytes, &encrypted_data)?;
+
+            let output_path = Path::new(&output_folder).join(package_name.clone() + ".bin");
+
+            fs::create_dir_all(&output_folder)?;
+            let mut out_file = OpenOptions::new()
+                .append(true)
+                .create(true)
+                .open(output_path)?;
+
+            out_file.write_all(&out_data[extra_segment_size..])?;
+
+            println!("-- Saved to file!");
+        }
+    }
+
+    Ok(())
+}