use std::collections::HashMap;
use std::io::{Cursor, Read};
use byteorder::{BigEndian, ReadBytesExt};
use tracing::{debug, trace};
use crate::utils::{read_cstring_from, read_uint40_from};
use crate::{Error, Result};
#[derive(Debug, Clone)]
pub struct DownloadHeader {
pub magic: [u8; 2],
pub version: u8,
pub ekey_size: u8,
pub has_checksum: bool,
pub entry_count: u32,
pub tag_count: u16,
pub flag_size: u8,
pub base_priority: i8,
pub unknown: u32,
}
impl DownloadHeader {
pub fn parse<R: Read>(reader: &mut R) -> Result<Self> {
let mut magic = [0u8; 2];
reader.read_exact(&mut magic)?;
if magic != [b'D', b'L'] {
return Err(Error::IOError(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("Invalid download manifest magic: {magic:?}"),
)));
}
let version = reader.read_u8()?;
let ekey_size = reader.read_u8()?;
let has_checksum = reader.read_u8()? != 0;
let entry_count = reader.read_u32::<BigEndian>()?;
let tag_count = reader.read_u16::<BigEndian>()?;
let mut flag_size = 0;
let mut base_priority = 0i8;
let mut unknown = 0u32;
if version >= 2 {
flag_size = reader.read_u8()?;
if version >= 3 {
base_priority = reader.read_i8()?;
let mut bytes = [0u8; 3];
reader.read_exact(&mut bytes)?;
unknown = u32::from_be_bytes([0, bytes[0], bytes[1], bytes[2]]);
}
}
Ok(DownloadHeader {
magic,
version,
ekey_size,
has_checksum,
entry_count,
tag_count,
flag_size,
base_priority,
unknown,
})
}
}
#[derive(Debug, Clone)]
pub struct DownloadEntry {
pub ekey: Vec<u8>,
pub compressed_size: u64,
pub priority: i8,
pub checksum: Option<u32>,
pub flags: Vec<u8>,
}
impl DownloadEntry {
pub fn parse<R: Read>(reader: &mut R, header: &DownloadHeader) -> Result<Self> {
let mut ekey = vec![0u8; header.ekey_size as usize];
reader.read_exact(&mut ekey)?;
let compressed_size = read_uint40_from(reader)?;
let raw_priority = reader.read_i8()?;
let priority = raw_priority - header.base_priority;
let checksum = if header.has_checksum {
Some(reader.read_u32::<BigEndian>()?)
} else {
None
};
let mut flags = vec![];
if header.version >= 2 && header.flag_size > 0 {
flags = vec![0u8; header.flag_size as usize];
reader.read_exact(&mut flags)?;
}
Ok(DownloadEntry {
ekey,
compressed_size,
priority,
checksum,
flags,
})
}
}
#[derive(Debug, Clone)]
pub struct DownloadTag {
pub name: String,
pub tag_type: u16,
pub mask: Vec<u8>,
}
#[derive(Debug, Clone)]
pub struct DownloadManifest {
pub header: DownloadHeader,
pub entries: HashMap<Vec<u8>, DownloadEntry>,
pub priority_order: Vec<Vec<u8>>,
pub tags: Vec<DownloadTag>,
}
impl DownloadManifest {
pub fn parse(data: &[u8]) -> Result<Self> {
let mut cursor = Cursor::new(data);
let header = DownloadHeader::parse(&mut cursor)?;
debug!(
"Parsing download manifest v{} with {} entries and {} tags",
header.version, header.entry_count, header.tag_count
);
let mut entries = HashMap::with_capacity(header.entry_count as usize);
let mut priority_list = Vec::with_capacity(header.entry_count as usize);
for i in 0..header.entry_count {
let entry = DownloadEntry::parse(&mut cursor, &header)?;
trace!(
"Entry {}: EKey {:02x?} priority={} size={}",
i,
&entry.ekey[..4.min(entry.ekey.len())],
entry.priority,
entry.compressed_size
);
priority_list.push((entry.priority, entry.ekey.clone()));
entries.insert(entry.ekey.clone(), entry);
}
priority_list.sort_by_key(|(priority, _)| *priority);
let priority_order: Vec<Vec<u8>> =
priority_list.into_iter().map(|(_, ekey)| ekey).collect();
let mut tags = Vec::with_capacity(header.tag_count as usize);
let bytes_per_tag = header.entry_count.div_ceil(8) as usize;
for i in 0..header.tag_count {
let name = read_cstring_from(&mut cursor)?;
let tag_type = cursor.read_u16::<BigEndian>()?;
let mut mask = vec![0u8; bytes_per_tag];
cursor.read_exact(&mut mask)?;
trace!("Tag {}: '{}' type={}", i, name, tag_type);
tags.push(DownloadTag {
name,
tag_type,
mask,
});
}
debug!(
"Parsed {} entries with {} priority levels",
entries.len(),
entries
.values()
.map(|e| e.priority)
.collect::<std::collections::HashSet<_>>()
.len()
);
Ok(DownloadManifest {
header,
entries,
priority_order,
tags,
})
}
pub fn get_priority_files(&self, max_priority: i8) -> Vec<&DownloadEntry> {
self.priority_order
.iter()
.filter_map(|ekey| {
let entry = self.entries.get(ekey)?;
if entry.priority <= max_priority {
Some(entry)
} else {
None
}
})
.collect()
}
pub fn get_files_for_tags(&self, tag_names: &[&str]) -> Vec<&DownloadEntry> {
let mut combined_mask = vec![0u8; self.header.entry_count.div_ceil(8) as usize];
for tag in &self.tags {
if tag_names.contains(&tag.name.as_str()) {
for (i, byte) in tag.mask.iter().enumerate() {
combined_mask[i] |= byte;
}
}
}
let mut result = Vec::new();
for (index, ekey) in self.priority_order.iter().enumerate() {
let byte_index = index / 8;
let bit_index = index % 8;
if byte_index < combined_mask.len() {
let bit = (combined_mask[byte_index] >> (7 - bit_index)) & 1;
if bit == 1 {
if let Some(entry) = self.entries.get(ekey) {
result.push(entry);
}
}
}
}
result
}
pub fn get_download_size(&self, max_priority: i8) -> u64 {
self.get_priority_files(max_priority)
.iter()
.map(|e| e.compressed_size)
.sum()
}
pub fn get_essential_files(&self) -> Vec<&DownloadEntry> {
self.get_priority_files(0)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_download_header_v1() {
let data = vec![
b'D', b'L', 1, 16, 0, 0, 0, 0, 2, 0, 1, ];
let mut cursor = Cursor::new(data);
let header = DownloadHeader::parse(&mut cursor).unwrap();
assert_eq!(header.magic, [b'D', b'L']);
assert_eq!(header.version, 1);
assert_eq!(header.ekey_size, 16);
assert!(!header.has_checksum);
assert_eq!(header.entry_count, 2);
assert_eq!(header.tag_count, 1);
assert_eq!(header.flag_size, 0); }
#[test]
fn test_download_header_v3() {
let data = vec![
b'D', b'L', 3, 16, 1, 0, 0, 0, 10, 0, 3, 2, 254u8, 0, 0, 0, ];
let mut cursor = Cursor::new(data);
let header = DownloadHeader::parse(&mut cursor).unwrap();
assert_eq!(header.version, 3);
assert!(header.has_checksum);
assert_eq!(header.entry_count, 10);
assert_eq!(header.tag_count, 3);
assert_eq!(header.flag_size, 2);
assert_eq!(header.base_priority, -2);
}
#[test]
fn test_priority_sorting() {
let mut entries = HashMap::new();
let entry1 = DownloadEntry {
ekey: vec![1; 16],
compressed_size: 1000,
priority: 2, checksum: None,
flags: vec![],
};
let entry2 = DownloadEntry {
ekey: vec![2; 16],
compressed_size: 2000,
priority: 0, checksum: None,
flags: vec![],
};
let entry3 = DownloadEntry {
ekey: vec![3; 16],
compressed_size: 3000,
priority: 1, checksum: None,
flags: vec![],
};
entries.insert(entry1.ekey.clone(), entry1);
entries.insert(entry2.ekey.clone(), entry2);
entries.insert(entry3.ekey.clone(), entry3);
let mut priority_list = vec![(2, vec![1; 16]), (0, vec![2; 16]), (1, vec![3; 16])];
priority_list.sort_by_key(|(p, _)| *p);
let priority_order: Vec<Vec<u8>> =
priority_list.into_iter().map(|(_, ekey)| ekey).collect();
assert_eq!(priority_order[0], vec![2; 16]); assert_eq!(priority_order[1], vec![3; 16]); assert_eq!(priority_order[2], vec![1; 16]); }
}