use byteorder::{BigEndian, ReadBytesExt};
use std::collections::HashMap;
use std::io::{Cursor, Read};
use tracing::{debug, trace, warn};
use crate::{Error, Result};
const ENCODING_MAGIC: [u8; 2] = [0x45, 0x4E];
#[derive(Debug, Clone)]
pub struct EncodingHeader {
pub magic: [u8; 2],
pub version: u8,
pub ckey_hash_size: u8,
pub ekey_hash_size: u8,
pub ckey_page_size_kb: u16,
pub ekey_page_size_kb: u16,
pub ckey_page_count: u32,
pub ekey_page_count: u32,
pub unk: u8,
pub espec_block_size: u32,
}
#[derive(Debug, Clone)]
pub struct PageInfo {
pub first_hash: Vec<u8>,
pub checksum: [u8; 16],
}
#[derive(Debug, Clone)]
pub struct EncodingEntry {
pub content_key: Vec<u8>,
pub encoding_keys: Vec<Vec<u8>>,
pub size: u64,
}
pub struct EncodingFile {
pub header: EncodingHeader,
ckey_entries: HashMap<Vec<u8>, EncodingEntry>,
ekey_to_ckey: HashMap<Vec<u8>, Vec<u8>>,
}
impl EncodingFile {
pub fn parse(data: &[u8]) -> Result<Self> {
let mut cursor = Cursor::new(data);
let header = Self::parse_header(&mut cursor)?;
debug!(
"Parsed encoding header: version={}, ckey_pages={}, ekey_pages={}, ckey_page_size_kb={}, ekey_page_size_kb={}, espec_table_size={}",
header.version,
header.ckey_page_count,
header.ekey_page_count,
header.ckey_page_size_kb,
header.ekey_page_size_kb,
header.espec_block_size
);
let mut espec_data = vec![0u8; header.espec_block_size as usize];
cursor.read_exact(&mut espec_data)?;
debug!("Read ESpec string table: {} bytes", espec_data.len());
let ckey_page_table = Self::parse_page_table(
&mut cursor,
header.ckey_page_count as usize,
header.ckey_hash_size as usize,
)?;
trace!("Parsed {} CKey page table entries", ckey_page_table.len());
let _ekey_page_table = Self::parse_page_table(
&mut cursor,
header.ekey_page_count as usize,
header.ekey_hash_size as usize,
)?;
trace!("Parsed {} EKey page table entries", _ekey_page_table.len());
let mut ckey_entries = HashMap::new();
let page_size = header.ckey_page_size_kb as usize * 1024;
let remaining_data = {
let current_pos = cursor.position() as usize;
&data[current_pos..]
};
let mut data_offset = 0;
for (i, page_info) in ckey_page_table.iter().enumerate() {
if data_offset + page_size <= remaining_data.len() {
let page_slice = &remaining_data[data_offset..data_offset + page_size];
let checksum = ::md5::compute(page_slice);
if checksum.as_ref() != page_info.checksum {
debug!(
"CKey page {} checksum mismatch (expected: {:?}, got: {:?})",
i,
hex::encode(page_info.checksum),
hex::encode(checksum.as_ref())
);
}
Self::parse_ckey_page(
page_slice,
header.ckey_hash_size,
header.ekey_hash_size,
&mut ckey_entries,
)?;
}
data_offset += page_size;
}
cursor.set_position(cursor.position() + (header.ckey_page_count as u64 * page_size as u64));
debug!("Parsed {} CKey entries", ckey_entries.len());
let mut ekey_to_ckey = HashMap::new();
for entry in ckey_entries.values() {
for ekey in &entry.encoding_keys {
ekey_to_ckey.insert(ekey.clone(), entry.content_key.clone());
}
}
debug!(
"Built EKey→CKey reverse mapping with {} entries",
ekey_to_ckey.len()
);
Ok(Self {
header,
ckey_entries,
ekey_to_ckey,
})
}
fn parse_header<R: Read>(reader: &mut R) -> Result<EncodingHeader> {
let mut magic = [0u8; 2];
reader.read_exact(&mut magic)?;
if magic != ENCODING_MAGIC {
return Err(Error::BadMagic);
}
let version = reader.read_u8()?;
if version != 1 {
warn!("Unexpected encoding version: {}", version);
}
let ckey_hash_size = reader.read_u8()?;
let ekey_hash_size = reader.read_u8()?;
let ckey_page_size_kb = reader.read_u16::<BigEndian>()?; let ekey_page_size_kb = reader.read_u16::<BigEndian>()?; let ckey_page_count = reader.read_u32::<BigEndian>()?; let ekey_page_count = reader.read_u32::<BigEndian>()?; let unk = reader.read_u8()?;
let espec_block_size = reader.read_u32::<BigEndian>()?;
Ok(EncodingHeader {
magic,
version,
ckey_hash_size,
ekey_hash_size,
ckey_page_size_kb,
ekey_page_size_kb,
ckey_page_count,
ekey_page_count,
unk,
espec_block_size,
})
}
fn parse_page_table<R: Read>(
reader: &mut R,
page_count: usize,
hash_size: usize,
) -> Result<Vec<PageInfo>> {
let mut pages = Vec::with_capacity(page_count);
for _ in 0..page_count {
let mut first_hash = vec![0u8; hash_size];
reader.read_exact(&mut first_hash)?;
let mut checksum = [0u8; 16];
reader.read_exact(&mut checksum)?;
pages.push(PageInfo {
first_hash,
checksum,
});
}
Ok(pages)
}
fn parse_ckey_page(
data: &[u8],
ckey_size: u8,
ekey_size: u8,
entries: &mut HashMap<Vec<u8>, EncodingEntry>,
) -> Result<()> {
let mut offset = 0;
while offset < data.len() {
if offset + 6 > data.len() || data[offset..].iter().all(|&b| b == 0) {
break;
}
let key_count = data[offset];
offset += 1;
if key_count == 0 {
break; }
if offset + 5 > data.len() {
break;
}
let size = crate::utils::read_uint40_be(&data[offset..offset + 5])?;
offset += 5;
if offset + ckey_size as usize > data.len() {
break;
}
let ckey = data[offset..offset + ckey_size as usize].to_vec();
offset += ckey_size as usize;
let mut ekeys = Vec::new();
for _ in 0..key_count {
if offset + ekey_size as usize > data.len() {
break;
}
let ekey = data[offset..offset + ekey_size as usize].to_vec();
offset += ekey_size as usize;
ekeys.push(ekey);
}
entries.insert(
ckey.clone(),
EncodingEntry {
content_key: ckey,
encoding_keys: ekeys,
size,
},
);
}
Ok(())
}
pub fn lookup_by_ckey(&self, ckey: &[u8]) -> Option<&EncodingEntry> {
self.ckey_entries.get(ckey)
}
pub fn lookup_by_ekey(&self, ekey: &[u8]) -> Option<&Vec<u8>> {
self.ekey_to_ckey.get(ekey)
}
pub fn get_ekey_for_ckey(&self, ckey: &[u8]) -> Option<&Vec<u8>> {
self.ckey_entries
.get(ckey)
.and_then(|entry| entry.encoding_keys.first())
}
pub fn get_file_size(&self, ckey: &[u8]) -> Option<u64> {
self.ckey_entries.get(ckey).map(|entry| entry.size)
}
pub fn ckey_count(&self) -> usize {
self.ckey_entries.len()
}
pub fn ekey_count(&self) -> usize {
self.ekey_to_ckey.len()
}
pub fn get_sample_ckeys(&self, limit: usize) -> Vec<String> {
self.ckey_entries
.keys()
.take(limit)
.map(hex::encode)
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_encoding_header_size() {
let header_size = 2 + 1 + 1 + 1 + 2 + 2 + 4 + 4 + 1 + 4;
assert_eq!(header_size, 22);
}
#[test]
fn test_parse_empty_encoding() {
let mut data = Vec::new();
data.extend_from_slice(&ENCODING_MAGIC);
data.push(1);
data.push(16); data.push(16); data.extend_from_slice(&0u16.to_be_bytes()); data.extend_from_slice(&0u16.to_be_bytes()); data.extend_from_slice(&0u32.to_be_bytes()); data.extend_from_slice(&0u32.to_be_bytes()); data.push(0);
data.extend_from_slice(&0u32.to_be_bytes());
let result = EncodingFile::parse(&data);
assert!(result.is_ok());
let encoding = result.unwrap();
assert_eq!(encoding.header.version, 1);
assert_eq!(encoding.header.ckey_hash_size, 16);
assert_eq!(encoding.header.ekey_hash_size, 16);
assert_eq!(encoding.ckey_count(), 0);
assert_eq!(encoding.ekey_count(), 0);
}
#[test]
fn test_invalid_magic() {
let mut data = vec![0xFF, 0xFF]; data.push(1);
let result = EncodingFile::parse(&data);
assert!(matches!(result, Err(Error::BadMagic)));
}
}