1use std::collections::HashMap;
7use std::io::{Cursor, Read};
8
9use byteorder::{BigEndian, ReadBytesExt};
10use tracing::{debug, trace};
11
12use crate::utils::{read_cstring_from, read_uint40_from};
13use crate::{Error, Result};
14
15#[derive(Debug, Clone)]
17pub struct DownloadHeader {
18 pub magic: [u8; 2],
20 pub version: u8,
22 pub ekey_size: u8,
24 pub has_checksum: bool,
26 pub entry_count: u32,
28 pub tag_count: u16,
30 pub flag_size: u8,
32 pub base_priority: i8,
34 pub unknown: u32,
36}
37
38impl DownloadHeader {
39 pub fn parse<R: Read>(reader: &mut R) -> Result<Self> {
41 let mut magic = [0u8; 2];
42 reader.read_exact(&mut magic)?;
43
44 if magic != [b'D', b'L'] {
45 return Err(Error::IOError(std::io::Error::new(
46 std::io::ErrorKind::InvalidData,
47 format!("Invalid download manifest magic: {magic:?}"),
48 )));
49 }
50
51 let version = reader.read_u8()?;
52 let ekey_size = reader.read_u8()?;
53 let has_checksum = reader.read_u8()? != 0;
54 let entry_count = reader.read_u32::<BigEndian>()?;
55 let tag_count = reader.read_u16::<BigEndian>()?;
56
57 let mut flag_size = 0;
58 let mut base_priority = 0i8;
59 let mut unknown = 0u32;
60
61 if version >= 2 {
62 flag_size = reader.read_u8()?;
63
64 if version >= 3 {
65 base_priority = reader.read_i8()?;
66 let mut bytes = [0u8; 3];
68 reader.read_exact(&mut bytes)?;
69 unknown = u32::from_be_bytes([0, bytes[0], bytes[1], bytes[2]]);
70 }
71 }
72
73 Ok(DownloadHeader {
74 magic,
75 version,
76 ekey_size,
77 has_checksum,
78 entry_count,
79 tag_count,
80 flag_size,
81 base_priority,
82 unknown,
83 })
84 }
85}
86
87#[derive(Debug, Clone)]
89pub struct DownloadEntry {
90 pub ekey: Vec<u8>,
92 pub compressed_size: u64,
94 pub priority: i8,
96 pub checksum: Option<u32>,
98 pub flags: Vec<u8>,
100}
101
102impl DownloadEntry {
103 pub fn parse<R: Read>(reader: &mut R, header: &DownloadHeader) -> Result<Self> {
105 let mut ekey = vec![0u8; header.ekey_size as usize];
106 reader.read_exact(&mut ekey)?;
107
108 let compressed_size = read_uint40_from(reader)?;
109
110 let raw_priority = reader.read_i8()?;
112 let priority = raw_priority - header.base_priority;
113
114 let checksum = if header.has_checksum {
115 Some(reader.read_u32::<BigEndian>()?)
116 } else {
117 None
118 };
119
120 let mut flags = vec![];
121 if header.version >= 2 && header.flag_size > 0 {
122 flags = vec![0u8; header.flag_size as usize];
123 reader.read_exact(&mut flags)?;
124 }
125
126 Ok(DownloadEntry {
127 ekey,
128 compressed_size,
129 priority,
130 checksum,
131 flags,
132 })
133 }
134}
135
136#[derive(Debug, Clone)]
138pub struct DownloadTag {
139 pub name: String,
141 pub tag_type: u16,
143 pub mask: Vec<u8>,
145}
146
147#[derive(Debug, Clone)]
149pub struct DownloadManifest {
150 pub header: DownloadHeader,
152 pub entries: HashMap<Vec<u8>, DownloadEntry>,
154 pub priority_order: Vec<Vec<u8>>,
156 pub tags: Vec<DownloadTag>,
158}
159
160impl DownloadManifest {
161 pub fn parse(data: &[u8]) -> Result<Self> {
163 let mut cursor = Cursor::new(data);
164
165 let header = DownloadHeader::parse(&mut cursor)?;
167
168 debug!(
169 "Parsing download manifest v{} with {} entries and {} tags",
170 header.version, header.entry_count, header.tag_count
171 );
172
173 let mut entries = HashMap::with_capacity(header.entry_count as usize);
175 let mut priority_list = Vec::with_capacity(header.entry_count as usize);
176
177 for i in 0..header.entry_count {
178 let entry = DownloadEntry::parse(&mut cursor, &header)?;
179 trace!(
180 "Entry {}: EKey {:02x?} priority={} size={}",
181 i,
182 &entry.ekey[..4.min(entry.ekey.len())],
183 entry.priority,
184 entry.compressed_size
185 );
186 priority_list.push((entry.priority, entry.ekey.clone()));
187 entries.insert(entry.ekey.clone(), entry);
188 }
189
190 priority_list.sort_by_key(|(priority, _)| *priority);
192 let priority_order: Vec<Vec<u8>> =
193 priority_list.into_iter().map(|(_, ekey)| ekey).collect();
194
195 let mut tags = Vec::with_capacity(header.tag_count as usize);
197 let bytes_per_tag = header.entry_count.div_ceil(8) as usize;
198
199 for i in 0..header.tag_count {
200 let name = read_cstring_from(&mut cursor)?;
201 let tag_type = cursor.read_u16::<BigEndian>()?;
202
203 let mut mask = vec![0u8; bytes_per_tag];
204 cursor.read_exact(&mut mask)?;
205
206 trace!("Tag {}: '{}' type={}", i, name, tag_type);
207
208 tags.push(DownloadTag {
209 name,
210 tag_type,
211 mask,
212 });
213 }
214
215 debug!(
216 "Parsed {} entries with {} priority levels",
217 entries.len(),
218 entries
219 .values()
220 .map(|e| e.priority)
221 .collect::<std::collections::HashSet<_>>()
222 .len()
223 );
224
225 Ok(DownloadManifest {
226 header,
227 entries,
228 priority_order,
229 tags,
230 })
231 }
232
233 pub fn get_priority_files(&self, max_priority: i8) -> Vec<&DownloadEntry> {
235 self.priority_order
236 .iter()
237 .filter_map(|ekey| {
238 let entry = self.entries.get(ekey)?;
239 if entry.priority <= max_priority {
240 Some(entry)
241 } else {
242 None
243 }
244 })
245 .collect()
246 }
247
248 pub fn get_files_for_tags(&self, tag_names: &[&str]) -> Vec<&DownloadEntry> {
250 let mut combined_mask = vec![0u8; self.header.entry_count.div_ceil(8) as usize];
252
253 for tag in &self.tags {
254 if tag_names.contains(&tag.name.as_str()) {
255 for (i, byte) in tag.mask.iter().enumerate() {
257 combined_mask[i] |= byte;
258 }
259 }
260 }
261
262 let mut result = Vec::new();
264 for (index, ekey) in self.priority_order.iter().enumerate() {
265 let byte_index = index / 8;
266 let bit_index = index % 8;
267
268 if byte_index < combined_mask.len() {
269 let bit = (combined_mask[byte_index] >> (7 - bit_index)) & 1;
270 if bit == 1 {
271 if let Some(entry) = self.entries.get(ekey) {
272 result.push(entry);
273 }
274 }
275 }
276 }
277
278 result
279 }
280
281 pub fn get_download_size(&self, max_priority: i8) -> u64 {
283 self.get_priority_files(max_priority)
284 .iter()
285 .map(|e| e.compressed_size)
286 .sum()
287 }
288
289 pub fn get_essential_files(&self) -> Vec<&DownloadEntry> {
291 self.get_priority_files(0)
292 }
293}
294
295#[cfg(test)]
296mod tests {
297 use super::*;
298
299 #[test]
300 fn test_download_header_v1() {
301 let data = vec![
302 b'D', b'L', 1, 16, 0, 0, 0, 0, 2, 0, 1, ];
309
310 let mut cursor = Cursor::new(data);
311 let header = DownloadHeader::parse(&mut cursor).unwrap();
312
313 assert_eq!(header.magic, [b'D', b'L']);
314 assert_eq!(header.version, 1);
315 assert_eq!(header.ekey_size, 16);
316 assert!(!header.has_checksum);
317 assert_eq!(header.entry_count, 2);
318 assert_eq!(header.tag_count, 1);
319 assert_eq!(header.flag_size, 0); }
321
322 #[test]
323 fn test_download_header_v3() {
324 let data = vec![
325 b'D', b'L', 3, 16, 1, 0, 0, 0, 10, 0, 3, 2, 254u8, 0, 0, 0, ];
335
336 let mut cursor = Cursor::new(data);
337 let header = DownloadHeader::parse(&mut cursor).unwrap();
338
339 assert_eq!(header.version, 3);
340 assert!(header.has_checksum);
341 assert_eq!(header.entry_count, 10);
342 assert_eq!(header.tag_count, 3);
343 assert_eq!(header.flag_size, 2);
344 assert_eq!(header.base_priority, -2);
345 }
346
347 #[test]
348 fn test_priority_sorting() {
349 let mut entries = HashMap::new();
351
352 let entry1 = DownloadEntry {
353 ekey: vec![1; 16],
354 compressed_size: 1000,
355 priority: 2, checksum: None,
357 flags: vec![],
358 };
359
360 let entry2 = DownloadEntry {
361 ekey: vec![2; 16],
362 compressed_size: 2000,
363 priority: 0, checksum: None,
365 flags: vec![],
366 };
367
368 let entry3 = DownloadEntry {
369 ekey: vec![3; 16],
370 compressed_size: 3000,
371 priority: 1, checksum: None,
373 flags: vec![],
374 };
375
376 entries.insert(entry1.ekey.clone(), entry1);
377 entries.insert(entry2.ekey.clone(), entry2);
378 entries.insert(entry3.ekey.clone(), entry3);
379
380 let mut priority_list = vec![(2, vec![1; 16]), (0, vec![2; 16]), (1, vec![3; 16])];
382 priority_list.sort_by_key(|(p, _)| *p);
383
384 let priority_order: Vec<Vec<u8>> =
385 priority_list.into_iter().map(|(_, ekey)| ekey).collect();
386
387 assert_eq!(priority_order[0], vec![2; 16]); assert_eq!(priority_order[1], vec![3; 16]); assert_eq!(priority_order[2], vec![1; 16]); }
392}