1use crate::{
4 Error, Result,
5 compression::{compress, flags as compression_flags},
6 crypto::{encrypt_block, hash_string, hash_type, het_hash, jenkins_hash},
7 header::{FormatVersion, MpqHeaderV4Data},
8 special_files::{AttributeFlags, Attributes, FileAttributes},
9 tables::{BetHeader, BlockEntry, BlockTable, HashEntry, HashTable, HetHeader, HiBlockTable},
10};
11use md5::{Digest, Md5};
12use std::fs::{self};
13use std::io::{Read, Seek, SeekFrom, Write};
14use std::path::{Path, PathBuf};
15use tempfile::NamedTempFile;
16
17trait WriteLittleEndian: Write {
19 fn write_u16_le(&mut self, value: u16) -> Result<()> {
20 self.write_all(&value.to_le_bytes())?;
21 Ok(())
22 }
23
24 fn write_u32_le(&mut self, value: u32) -> Result<()> {
25 self.write_all(&value.to_le_bytes())?;
26 Ok(())
27 }
28
29 fn write_u64_le(&mut self, value: u64) -> Result<()> {
30 self.write_all(&value.to_le_bytes())?;
31 Ok(())
32 }
33}
34
35impl<W: Write> WriteLittleEndian for W {}
36
37#[derive(Debug)]
39struct PendingFile {
40 source: FileSource,
42 archive_name: String,
44 compression: u8,
46 encrypt: bool,
48 use_fix_key: bool,
50 locale: u16,
52}
53
54#[derive(Debug)]
55enum FileSource {
56 Path(PathBuf),
57 Data(Vec<u8>),
58}
59
60struct FileWriteParams<'a> {
62 file_data: &'a [u8],
64 archive_name: &'a str,
66 compression: u8,
68 encrypt: bool,
70 use_fix_key: bool,
72 sector_size: usize,
74 file_pos: u64,
76}
77
78struct HeaderWriteParams {
80 archive_size: u64,
81 hash_table_pos: u64,
82 block_table_pos: u64,
83 hash_table_size: u32,
84 block_table_size: u32,
85 hi_block_table_pos: Option<u64>,
86 het_table_pos: Option<u64>,
87 bet_table_pos: Option<u64>,
88 _het_table_size: Option<u64>,
89 _bet_table_size: Option<u64>,
90 v4_data: Option<MpqHeaderV4Data>,
92}
93
94#[derive(Debug, Clone)]
96pub enum ListfileOption {
97 Generate,
99 External(PathBuf),
101 None,
103}
104
105#[derive(Debug, Clone)]
107pub enum AttributesOption {
108 GenerateCrc32,
110 GenerateFull,
112 External(PathBuf),
114 None,
116}
117
118#[derive(Debug)]
160pub struct ArchiveBuilder {
161 version: FormatVersion,
163 block_size: u16,
165 pending_files: Vec<PendingFile>,
167 listfile_option: ListfileOption,
169 attributes_option: AttributesOption,
171 default_compression: u8,
173 generate_crcs: bool,
175 compress_tables: bool,
177 table_compression: u8,
179}
180
181impl ArchiveBuilder {
182 pub fn new() -> Self {
184 Self {
185 version: FormatVersion::V1,
186 block_size: 5, pending_files: Vec::new(),
188 listfile_option: ListfileOption::Generate,
189 attributes_option: AttributesOption::None,
190 default_compression: compression_flags::ZLIB,
191 generate_crcs: false,
192 compress_tables: false, table_compression: compression_flags::ZLIB,
194 }
195 }
196
197 pub fn version(mut self, version: FormatVersion) -> Self {
199 self.version = version;
200 self
201 }
202
203 pub fn block_size(mut self, block_size: u16) -> Self {
225 self.block_size = block_size;
226 self
227 }
228
229 pub fn default_compression(mut self, compression: u8) -> Self {
231 self.default_compression = compression;
232 self
233 }
234
235 pub fn listfile_option(mut self, option: ListfileOption) -> Self {
237 self.listfile_option = option;
238 self
239 }
240
241 pub fn generate_crcs(mut self, generate: bool) -> Self {
263 self.generate_crcs = generate;
264 if generate && matches!(self.attributes_option, AttributesOption::None) {
266 self.attributes_option = AttributesOption::GenerateCrc32;
267 }
268 self
269 }
270
271 pub fn attributes_option(mut self, option: AttributesOption) -> Self {
290 if matches!(
292 option,
293 AttributesOption::GenerateCrc32 | AttributesOption::GenerateFull
294 ) {
295 self.generate_crcs = true;
296 }
297 self.attributes_option = option;
298 self
299 }
300
301 pub fn compress_tables(mut self, compress: bool) -> Self {
325 self.compress_tables = compress;
326 self
327 }
328
329 pub fn table_compression(mut self, compression: u8) -> Self {
352 self.table_compression = compression;
353 self
354 }
355
356 pub fn add_file<P: AsRef<Path>>(mut self, path: P, archive_name: &str) -> Self {
381 self.pending_files.push(PendingFile {
382 source: FileSource::Path(path.as_ref().to_path_buf()),
383 archive_name: crate::path::normalize_mpq_path(archive_name),
384 compression: self.default_compression,
385 encrypt: false,
386 use_fix_key: false,
387 locale: 0, });
389 self
390 }
391
392 pub fn add_file_with_options<P: AsRef<Path>>(
419 mut self,
420 path: P,
421 archive_name: &str,
422 compression: u8,
423 encrypt: bool,
424 locale: u16,
425 ) -> Self {
426 self.pending_files.push(PendingFile {
427 source: FileSource::Path(path.as_ref().to_path_buf()),
428 archive_name: crate::path::normalize_mpq_path(archive_name),
429 compression,
430 encrypt,
431 use_fix_key: false,
432 locale,
433 });
434 self
435 }
436
437 pub fn add_file_data(mut self, data: Vec<u8>, archive_name: &str) -> Self {
463 self.pending_files.push(PendingFile {
464 source: FileSource::Data(data),
465 archive_name: crate::path::normalize_mpq_path(archive_name),
466 compression: self.default_compression,
467 encrypt: false,
468 use_fix_key: false,
469 locale: 0,
470 });
471 self
472 }
473
474 pub fn add_file_data_with_options(
502 mut self,
503 data: Vec<u8>,
504 archive_name: &str,
505 compression: u8,
506 encrypt: bool,
507 locale: u16,
508 ) -> Self {
509 self.pending_files.push(PendingFile {
510 source: FileSource::Data(data),
511 archive_name: crate::path::normalize_mpq_path(archive_name),
512 compression,
513 encrypt,
514 use_fix_key: false,
515 locale,
516 });
517 self
518 }
519
520 pub fn add_file_with_encryption<P: AsRef<Path>>(
522 mut self,
523 path: P,
524 archive_name: &str,
525 compression: u8,
526 use_fix_key: bool,
527 locale: u16,
528 ) -> Self {
529 self.pending_files.push(PendingFile {
530 source: FileSource::Path(path.as_ref().to_path_buf()),
531 archive_name: crate::path::normalize_mpq_path(archive_name),
532 compression,
533 encrypt: true,
534 use_fix_key,
535 locale,
536 });
537 self
538 }
539
540 pub fn add_file_data_with_encryption(
542 mut self,
543 data: Vec<u8>,
544 archive_name: &str,
545 compression: u8,
546 use_fix_key: bool,
547 locale: u16,
548 ) -> Self {
549 self.pending_files.push(PendingFile {
550 source: FileSource::Data(data),
551 archive_name: crate::path::normalize_mpq_path(archive_name),
552 compression,
553 encrypt: true,
554 use_fix_key,
555 locale,
556 });
557 self
558 }
559
560 fn calculate_hash_table_size(&self) -> u32 {
562 let file_count = self.pending_files.len()
563 + match &self.listfile_option {
564 ListfileOption::Generate | ListfileOption::External(_) => 1,
565 ListfileOption::None => 0,
566 }
567 + match &self.attributes_option {
568 AttributesOption::GenerateCrc32
569 | AttributesOption::GenerateFull
570 | AttributesOption::External(_) => 1,
571 AttributesOption::None => 0,
572 };
573
574 let optimal_size = (file_count * 2).max(16) as u32;
576
577 optimal_size.next_power_of_two()
579 }
580
581 pub fn build<P: AsRef<Path>>(mut self, path: P) -> Result<()> {
583 let path = path.as_ref();
584
585 let mut temp_file = NamedTempFile::new_in(path.parent().unwrap_or_else(|| Path::new(".")))?;
587
588 self.prepare_listfile()?;
590
591 self.prepare_attributes()?;
593
594 {
596 let file = temp_file.as_file_mut();
597 use std::io::{Seek as _, Write as _};
598
599 if self.version >= FormatVersion::V3 {
602 let header_size = self.version.header_size() as usize;
605
606 let estimated_file_data_size: usize = self
608 .pending_files
609 .iter()
610 .map(|f| match &f.source {
611 FileSource::Data(data) => data.len() / 10 + 1000, FileSource::Path(_) => 100_000, })
614 .sum();
615
616 let estimated_table_size = self.pending_files.len() * 1000; let total_estimated_size =
620 header_size + estimated_file_data_size + estimated_table_size;
621
622 log::debug!(
623 "Pre-allocating buffer of {total_estimated_size} bytes for v3+ archive (header: {header_size}, estimated data: {estimated_file_data_size}, tables: {estimated_table_size})"
624 );
625
626 let mut vec = Vec::with_capacity(total_estimated_size);
627 vec.resize(header_size, 0u8);
628 let mut buffer = std::io::Cursor::new(vec);
629 buffer.seek(SeekFrom::Start(header_size as u64))?;
630
631 self.write_archive(&mut buffer)?;
632
633 file.write_all(buffer.get_ref())?;
635 file.flush()?;
636 } else {
637 self.write_archive(file)?;
639 file.flush()?;
640 }
641 }
642
643 temp_file.persist(path).map_err(|e| Error::Io(e.error))?;
645
646 Ok(())
647 }
648
649 fn prepare_listfile(&mut self) -> Result<()> {
651 match &self.listfile_option {
652 ListfileOption::Generate => {
653 let mut content = String::new();
655 for file in &self.pending_files {
656 content.push_str(&file.archive_name);
657 content.push('\r');
658 content.push('\n');
659 }
660
661 content.push_str("(listfile)\r\n");
663
664 if matches!(
666 self.attributes_option,
667 AttributesOption::GenerateCrc32 | AttributesOption::GenerateFull
668 ) {
669 content.push_str("(attributes)\r\n");
670 }
671
672 self.pending_files.push(PendingFile {
673 source: FileSource::Data(content.into_bytes()),
674 archive_name: "(listfile)".to_string(),
675 compression: self.default_compression,
676 encrypt: false,
677 use_fix_key: false,
678 locale: 0,
679 });
680 }
681 ListfileOption::External(path) => {
682 let data = fs::read(path)?;
684
685 self.pending_files.push(PendingFile {
686 source: FileSource::Data(data),
687 archive_name: "(listfile)".to_string(),
688 compression: self.default_compression,
689 encrypt: false,
690 use_fix_key: false,
691 locale: 0,
692 });
693 }
694 ListfileOption::None => {}
695 }
696
697 Ok(())
698 }
699
700 fn prepare_attributes(&mut self) -> Result<()> {
702 match &self.attributes_option {
706 AttributesOption::GenerateCrc32 | AttributesOption::GenerateFull => {
707 }
711 AttributesOption::External(path) => {
712 let content = fs::read(path)?;
714 self.pending_files.push(PendingFile {
715 source: FileSource::Data(content),
716 archive_name: "(attributes)".to_string(),
717 compression: 0, encrypt: false,
719 use_fix_key: false,
720 locale: 0,
721 });
722 }
723 AttributesOption::None => {}
724 }
725
726 Ok(())
727 }
728
729 fn write_archive<W: Write + Seek + Read>(&self, writer: &mut W) -> Result<()> {
731 let use_het_bet = self.version >= FormatVersion::V3;
733
734 if use_het_bet {
735 return self.write_archive_with_het_bet(writer);
736 }
737
738 let hash_table_size = self.calculate_hash_table_size();
739 let mut block_table_size = self.pending_files.len() as u32;
740
741 if matches!(
743 self.attributes_option,
744 AttributesOption::GenerateCrc32 | AttributesOption::GenerateFull
745 ) {
746 block_table_size += 1;
747 }
748
749 let sector_size = crate::calculate_sector_size(self.block_size);
751
752 let header_size = self.version.header_size();
754 writer.seek(SeekFrom::Start(header_size as u64))?;
755
756 let mut hash_table = HashTable::new(hash_table_size as usize)?;
758 let mut block_table = BlockTable::new(block_table_size as usize)?;
759 let mut hi_block_table = if self.version >= FormatVersion::V2 {
760 Some(HiBlockTable::new(block_table_size as usize))
761 } else {
762 None
763 };
764
765 let collect_attributes = matches!(
767 self.attributes_option,
768 AttributesOption::GenerateCrc32 | AttributesOption::GenerateFull
769 );
770 let mut collected_attributes = if collect_attributes {
771 Some(Vec::new())
772 } else {
773 None
774 };
775
776 let mut actual_block_index = 0;
778 for pending_file in self.pending_files.iter() {
779 if pending_file.archive_name == "(attributes)" && collect_attributes {
781 continue;
782 }
783
784 let file_pos = writer.stream_position()?;
785
786 let file_data = match &pending_file.source {
788 FileSource::Path(path) => fs::read(path)?,
789 FileSource::Data(data) => data.clone(),
790 };
791
792 let params = FileWriteParams {
794 file_data: &file_data,
795 archive_name: &pending_file.archive_name,
796 compression: pending_file.compression,
797 encrypt: pending_file.encrypt,
798 use_fix_key: pending_file.use_fix_key,
799 sector_size,
800 file_pos,
801 };
802
803 let (compressed_size, flags, file_attr) = if collect_attributes {
804 self.write_file_with_attributes(writer, ¶ms)?
805 } else {
806 let (size, flags) = self.write_file(writer, ¶ms)?;
807 (size, flags, FileAttributes::new())
808 };
809
810 if let Some(ref mut attrs) = collected_attributes {
812 attrs.push(file_attr);
813 }
814
815 self.add_to_hash_table(
817 &mut hash_table,
818 &pending_file.archive_name,
819 actual_block_index as u32,
820 pending_file.locale,
821 )?;
822
823 let block_entry = BlockEntry {
825 file_pos: file_pos as u32, compressed_size: compressed_size as u32,
827 file_size: file_data.len() as u32,
828 flags: flags | BlockEntry::FLAG_EXISTS,
829 };
830
831 if let Some(ref mut hi_table) = hi_block_table {
833 let high_bits = (file_pos >> 32) as u16;
834 hi_table.set(actual_block_index, high_bits);
835 }
836
837 if let Some(entry) = block_table.get_mut(actual_block_index) {
839 *entry = block_entry;
840 } else {
841 return Err(Error::invalid_format("Block index out of bounds"));
842 }
843
844 actual_block_index += 1;
845 }
846
847 if let Some(attrs) = collected_attributes {
849 log::debug!("Writing attributes for {} files", attrs.len());
850 self.write_attributes_file(
851 writer,
852 &mut hash_table,
853 &mut block_table,
854 attrs,
855 actual_block_index,
856 )?;
857 }
858
859 let hash_table_pos = writer.stream_position()?;
861 self.write_hash_table(writer, &hash_table)?;
862
863 let block_table_pos = writer.stream_position()?;
865 self.write_block_table(writer, &block_table)?;
866
867 let hi_block_table_pos = if let Some(ref hi_table) = hi_block_table {
869 if hi_table.is_needed() {
870 let pos = writer.stream_position()?;
871 self.write_hi_block_table(writer, hi_table)?;
872 Some(pos)
873 } else {
874 None
875 }
876 } else {
877 None
878 };
879
880 let archive_size = writer.stream_position()?;
882
883 writer.seek(SeekFrom::Start(0))?;
885 let header_params = HeaderWriteParams {
886 archive_size,
887 hash_table_pos,
888 block_table_pos,
889 hash_table_size,
890 block_table_size,
891 hi_block_table_pos,
892 het_table_pos: None,
893 bet_table_pos: None,
894 _het_table_size: None,
895 _bet_table_size: None,
896 v4_data: None, };
898 self.write_header(writer, &header_params)?;
899
900 Ok(())
903 }
904
905 fn write_archive_with_het_bet<W: Write + Seek + Read>(&self, writer: &mut W) -> Result<()> {
907 let mut block_table_size = self.pending_files.len() as u32;
908
909 if matches!(
911 self.attributes_option,
912 AttributesOption::GenerateCrc32 | AttributesOption::GenerateFull
913 ) {
914 block_table_size += 1;
915 }
916
917 let sector_size = crate::calculate_sector_size(self.block_size);
919
920 let header_size = self.version.header_size();
922 writer.seek(SeekFrom::Start(header_size as u64))?;
923
924 let mut block_table = BlockTable::new(block_table_size as usize)?;
926 let mut hi_block_table = Some(HiBlockTable::new(block_table_size as usize));
927
928 let collect_attributes = matches!(
930 self.attributes_option,
931 AttributesOption::GenerateCrc32 | AttributesOption::GenerateFull
932 );
933 let mut collected_attributes = if collect_attributes {
934 Some(Vec::new())
935 } else {
936 None
937 };
938
939 let mut actual_block_index = 0;
941 for pending_file in self.pending_files.iter() {
942 if pending_file.archive_name == "(attributes)" && collect_attributes {
944 continue;
945 }
946
947 let file_pos = writer.stream_position()?;
948
949 let file_data = match &pending_file.source {
951 FileSource::Path(path) => fs::read(path)?,
952 FileSource::Data(data) => data.clone(),
953 };
954
955 let params = FileWriteParams {
957 file_data: &file_data,
958 archive_name: &pending_file.archive_name,
959 compression: pending_file.compression,
960 encrypt: pending_file.encrypt,
961 use_fix_key: pending_file.use_fix_key,
962 sector_size,
963 file_pos,
964 };
965
966 let (compressed_size, flags, file_attr) = if collect_attributes {
967 self.write_file_with_attributes(writer, ¶ms)?
968 } else {
969 let (size, flags) = self.write_file(writer, ¶ms)?;
970 (size, flags, FileAttributes::new())
971 };
972
973 if let Some(ref mut attrs) = collected_attributes {
975 attrs.push(file_attr);
976 }
977
978 let block_entry = BlockEntry {
980 file_pos: file_pos as u32, compressed_size: compressed_size as u32,
982 file_size: file_data.len() as u32,
983 flags: flags | BlockEntry::FLAG_EXISTS,
984 };
985
986 if let Some(ref mut hi_table) = hi_block_table {
988 let high_bits = (file_pos >> 32) as u16;
989 hi_table.set(actual_block_index, high_bits);
990 }
991
992 if let Some(entry) = block_table.get_mut(actual_block_index) {
994 *entry = block_entry;
995 } else {
996 return Err(Error::invalid_format("Block index out of bounds"));
997 }
998
999 actual_block_index += 1;
1000 }
1001
1002 let has_collected_attributes = collected_attributes.is_some();
1004
1005 let hash_table_size = self.calculate_hash_table_size();
1007 let mut hash_table = HashTable::new(hash_table_size as usize)?;
1008
1009 let mut hash_block_index = 0;
1011 for pending_file in self.pending_files.iter() {
1012 if pending_file.archive_name == "(attributes)" && has_collected_attributes {
1014 continue;
1015 }
1016
1017 self.add_to_hash_table(
1018 &mut hash_table,
1019 &pending_file.archive_name,
1020 hash_block_index as u32,
1021 pending_file.locale,
1022 )?;
1023
1024 hash_block_index += 1;
1025 }
1026
1027 if let Some(attrs) = collected_attributes {
1029 self.write_attributes_file(
1030 writer,
1031 &mut hash_table,
1032 &mut block_table,
1033 attrs,
1034 actual_block_index,
1035 )?;
1036 }
1037
1038 let het_table_pos = writer.stream_position()?;
1040 let (het_data, _het_header) = self.create_het_table_with_hash_table(&hash_table)?;
1041 let (het_table_size, het_table_md5) = self.write_het_table(writer, &het_data, true)?;
1042
1043 let bet_table_pos = writer.stream_position()?;
1045 let (bet_data, _bet_header) = self.create_bet_table(&block_table)?;
1046 let (bet_table_size, bet_table_md5) = self.write_bet_table(writer, &bet_data, true)?;
1047
1048 let hash_table_pos = writer.stream_position()?;
1050 let hash_table_md5 = self.write_hash_table(writer, &hash_table)?;
1051
1052 let block_table_pos = writer.stream_position()?;
1054 let block_table_md5 = self.write_block_table(writer, &block_table)?;
1055
1056 let (hi_block_table_pos, hi_block_table_md5) = if let Some(ref hi_table) = hi_block_table {
1058 if hi_table.is_needed() {
1059 let pos = writer.stream_position()?;
1060 let md5 = self.write_hi_block_table(writer, hi_table)?;
1061 (Some(pos), md5)
1062 } else {
1063 (None, [0u8; 16])
1064 }
1065 } else {
1066 (None, [0u8; 16])
1067 };
1068
1069 let archive_size = writer.stream_position()?;
1071
1072 let _archive_end_pos = writer.stream_position()?;
1074
1075 writer.seek(SeekFrom::Start(0))?;
1077
1078 let actual_file_count = block_table_size; let v4_data = if self.version == FormatVersion::V4 {
1081 Some(MpqHeaderV4Data {
1082 hash_table_size_64: hash_table_size as u64 * 16, block_table_size_64: actual_file_count as u64 * 16, hi_block_table_size_64: if let Some(ref hi_table) = hi_block_table {
1085 if hi_table.is_needed() {
1086 actual_file_count as u64 * 2 } else {
1088 0
1089 }
1090 } else {
1091 0
1092 },
1093 het_table_size_64: het_table_size,
1094 bet_table_size_64: bet_table_size,
1095 raw_chunk_size: 0x4000, md5_block_table: block_table_md5,
1097 md5_hash_table: hash_table_md5,
1098 md5_hi_block_table: hi_block_table_md5,
1099 md5_bet_table: bet_table_md5,
1100 md5_het_table: het_table_md5,
1101 md5_mpq_header: [0u8; 16], })
1103 } else {
1104 None
1105 };
1106
1107 let header_params = HeaderWriteParams {
1108 archive_size,
1109 hash_table_pos,
1110 block_table_pos,
1111 hash_table_size,
1112 block_table_size: actual_file_count,
1113 hi_block_table_pos,
1114 het_table_pos: Some(het_table_pos),
1115 bet_table_pos: Some(bet_table_pos),
1116 _het_table_size: Some(het_table_size),
1117 _bet_table_size: Some(bet_table_size),
1118 v4_data,
1119 };
1120
1121 self.write_header(writer, &header_params)?;
1123
1124 if self.version == FormatVersion::V4 {
1126 self.finalize_v4_header_md5(writer)?;
1127 }
1128
1129 Ok(())
1130 }
1131
1132 fn write_file_with_attributes<W: Write>(
1134 &self,
1135 writer: &mut W,
1136 params: &FileWriteParams<'_>,
1137 ) -> Result<(usize, u32, FileAttributes)> {
1138 let (size, flags) = self.write_file(writer, params)?;
1139
1140 let mut file_attr = FileAttributes::new();
1142
1143 if matches!(
1145 self.attributes_option,
1146 AttributesOption::GenerateCrc32 | AttributesOption::GenerateFull
1147 ) {
1148 let crc32 = crc32fast::hash(params.file_data);
1149 file_attr.crc32 = Some(crc32);
1150 }
1151
1152 if matches!(self.attributes_option, AttributesOption::GenerateFull) {
1154 let mut hasher = Md5::new();
1155 hasher.update(params.file_data);
1156 let md5_result = hasher.finalize();
1157 let mut md5_bytes = [0u8; 16];
1158 md5_bytes.copy_from_slice(&md5_result);
1159 file_attr.md5 = Some(md5_bytes);
1160 }
1161
1162 if matches!(self.attributes_option, AttributesOption::GenerateFull) {
1164 use std::time::{SystemTime, UNIX_EPOCH};
1166 let duration = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
1167 let unix_seconds = duration.as_secs();
1168 let windows_seconds = unix_seconds + 11644473600;
1170 let filetime = windows_seconds * 10_000_000; file_attr.filetime = Some(filetime);
1172 }
1173
1174 Ok((size, flags, file_attr))
1175 }
1176
1177 fn write_file<W: Write>(
1179 &self,
1180 writer: &mut W,
1181 params: &FileWriteParams<'_>,
1182 ) -> Result<(usize, u32)> {
1183 let FileWriteParams {
1184 file_data,
1185 archive_name,
1186 compression,
1187 encrypt,
1188 use_fix_key,
1189 sector_size,
1190 file_pos,
1191 } = params;
1192 let mut flags = 0u32;
1193
1194 let is_single_unit = file_data.len() <= *sector_size;
1196
1197 if is_single_unit {
1198 flags |= BlockEntry::FLAG_SINGLE_UNIT;
1199
1200 if self.generate_crcs {
1202 flags |= BlockEntry::FLAG_SECTOR_CRC;
1203 }
1204
1205 let compressed_data = if *compression != 0 && !file_data.is_empty() {
1207 log::debug!("Compressing {archive_name} with method 0x{compression:02X}");
1208 let compressed = compress(file_data, *compression)?;
1209
1210 if compressed != *file_data {
1213 log::debug!(
1215 "Compression successful: {} -> {} bytes (including method byte)",
1216 file_data.len(),
1217 compressed.len()
1218 );
1219 flags |= BlockEntry::FLAG_COMPRESS;
1220 compressed
1221 } else {
1222 log::debug!("Compression not beneficial, storing uncompressed");
1224 file_data.to_vec()
1225 }
1226 } else {
1227 file_data.to_vec()
1228 };
1229
1230 let final_data = if *encrypt {
1232 flags |= BlockEntry::FLAG_ENCRYPTED;
1233 if *use_fix_key {
1234 flags |= BlockEntry::FLAG_FIX_KEY;
1235 }
1236 let key =
1237 self.calculate_file_key(archive_name, *file_pos, file_data.len() as u32, flags);
1238 let mut encrypted = compressed_data;
1239 self.encrypt_data(&mut encrypted, key);
1240 encrypted
1241 } else {
1242 compressed_data
1243 };
1244
1245 writer.write_all(&final_data)?;
1247
1248 if self.generate_crcs {
1250 let crc = adler2::adler32_slice(file_data);
1252 writer.write_u32_le(crc)?;
1253 log::debug!("Generated CRC for single unit file {archive_name}: 0x{crc:08X}");
1254 }
1255
1256 Ok((final_data.len(), flags))
1258 } else {
1259 let sector_count = file_data.len().div_ceil(*sector_size);
1261
1262 if self.generate_crcs {
1264 flags |= BlockEntry::FLAG_SECTOR_CRC;
1265 }
1266
1267 let offset_table_size = (sector_count + 1) * 4;
1269 let crc_table_size = if self.generate_crcs {
1270 sector_count * 4
1271 } else {
1272 0
1273 };
1274 let data_start = offset_table_size + crc_table_size;
1275
1276 let mut sector_offsets = vec![0u32; sector_count + 1];
1277 let mut sector_data = Vec::new();
1278 let mut sector_crcs = if self.generate_crcs {
1279 Vec::with_capacity(sector_count)
1280 } else {
1281 Vec::new()
1282 };
1283
1284 for (i, offset) in sector_offsets.iter_mut().enumerate().take(sector_count) {
1286 let sector_start = i * *sector_size;
1287 let sector_end = ((i + 1) * *sector_size).min(file_data.len());
1288 let sector_bytes = &file_data[sector_start..sector_end];
1289
1290 *offset = (data_start + sector_data.len()) as u32;
1291
1292 if self.generate_crcs {
1294 let crc = adler2::adler32_slice(sector_bytes);
1296 sector_crcs.push(crc);
1297 }
1298
1299 let compressed_sector = if *compression != 0 && !sector_bytes.is_empty() {
1301 let compressed = compress(sector_bytes, *compression)?;
1304 if compressed != *sector_bytes {
1305 flags |= BlockEntry::FLAG_COMPRESS;
1307 compressed
1308 } else {
1309 sector_bytes.to_vec()
1311 }
1312 } else {
1313 sector_bytes.to_vec()
1314 };
1315
1316 sector_data.extend_from_slice(&compressed_sector);
1317 }
1318
1319 sector_offsets[sector_count] = (data_start + sector_data.len()) as u32;
1321
1322 if self.generate_crcs {
1324 log::debug!(
1325 "Generated {} sector CRCs for file {}, first few: {:?}",
1326 sector_count,
1327 archive_name,
1328 §or_crcs[..5.min(sector_crcs.len())]
1329 );
1330 }
1331
1332 if *encrypt {
1334 flags |= BlockEntry::FLAG_ENCRYPTED;
1335 if *use_fix_key {
1336 flags |= BlockEntry::FLAG_FIX_KEY;
1337 }
1338 let key =
1339 self.calculate_file_key(archive_name, *file_pos, file_data.len() as u32, flags);
1340
1341 let original_offsets = sector_offsets.clone();
1343
1344 let offset_key = key.wrapping_sub(1);
1346 self.encrypt_data_u32(&mut sector_offsets, offset_key);
1347
1348 let mut encrypted_sectors = Vec::new();
1350 for (i, offset_pair) in original_offsets.windows(2).enumerate() {
1351 let start = (offset_pair[0] - data_start as u32) as usize;
1352 let end = (offset_pair[1] - data_start as u32) as usize;
1353
1354 let mut sector = sector_data[start..end].to_vec();
1355 let sector_key = key.wrapping_add(i as u32);
1356 self.encrypt_data(&mut sector, sector_key);
1357 encrypted_sectors.extend_from_slice(§or);
1358 }
1359
1360 sector_data = encrypted_sectors;
1361 }
1362
1363 for offset in §or_offsets {
1365 writer.write_u32_le(*offset)?;
1366 }
1367
1368 if self.generate_crcs {
1370 for crc in §or_crcs {
1371 writer.write_u32_le(*crc)?;
1372 }
1373 }
1374
1375 writer.write_all(§or_data)?;
1377
1378 let total_size = offset_table_size + sector_data.len();
1380 Ok((total_size, flags))
1381 }
1382 }
1383
1384 fn write_attributes_file<W: Write + Seek>(
1386 &self,
1387 writer: &mut W,
1388 hash_table: &mut HashTable,
1389 block_table: &mut BlockTable,
1390 file_attributes: Vec<FileAttributes>,
1391 block_index: usize,
1392 ) -> Result<()> {
1393 let flags = match self.attributes_option {
1395 AttributesOption::GenerateCrc32 => AttributeFlags::CRC32,
1396 AttributesOption::GenerateFull => {
1397 AttributeFlags::CRC32 | AttributeFlags::MD5 | AttributeFlags::FILETIME
1398 }
1399 _ => return Ok(()), };
1401
1402 let attributes = Attributes {
1404 version: Attributes::EXPECTED_VERSION,
1405 flags: AttributeFlags::new(flags),
1406 file_attributes,
1407 crc32: None, md5: None, filetime: None, };
1411
1412 let attributes_data = attributes.to_bytes()?;
1414
1415 let file_pos = writer.stream_position()?;
1417 writer.write_all(&attributes_data)?;
1418
1419 let block_entry = BlockEntry {
1422 file_pos: file_pos as u32,
1423 compressed_size: attributes_data.len() as u32,
1424 file_size: attributes_data.len() as u32,
1425 flags: BlockEntry::FLAG_EXISTS,
1426 };
1427
1428 if let Some(entry) = block_table.get_mut(block_index) {
1430 *entry = block_entry;
1431 } else {
1432 return Err(Error::invalid_format(
1433 "Invalid block index for attributes file",
1434 ));
1435 }
1436
1437 self.add_to_hash_table(hash_table, "(attributes)", block_index as u32, 0)?;
1439
1440 Ok(())
1441 }
1442
1443 fn add_to_hash_table(
1445 &self,
1446 hash_table: &mut HashTable,
1447 filename: &str,
1448 block_index: u32,
1449 locale: u16,
1450 ) -> Result<()> {
1451 let table_offset = hash_string(filename, hash_type::TABLE_OFFSET);
1452 let name_a = hash_string(filename, hash_type::NAME_A);
1453 let name_b = hash_string(filename, hash_type::NAME_B);
1454
1455 let table_size = hash_table.size() as u32;
1456 let mut index = table_offset & (table_size - 1);
1457
1458 loop {
1460 let entry = hash_table
1461 .get_mut(index as usize)
1462 .ok_or_else(|| Error::invalid_format("Hash table index out of bounds"))?;
1463
1464 if entry.is_empty() {
1465 *entry = HashEntry {
1467 name_1: name_a,
1468 name_2: name_b,
1469 locale,
1470 platform: 0, block_index,
1472 };
1473 break;
1474 }
1475
1476 if entry.name_1 == name_a && entry.name_2 == name_b && entry.locale == locale {
1478 return Err(Error::invalid_format(format!(
1479 "Duplicate file in archive: {filename}"
1480 )));
1481 }
1482
1483 index = (index + 1) & (table_size - 1);
1485 }
1486
1487 Ok(())
1488 }
1489
1490 fn write_hash_table<W: Write>(
1492 &self,
1493 writer: &mut W,
1494 hash_table: &HashTable,
1495 ) -> Result<[u8; 16]> {
1496 let mut table_data = Vec::new();
1498 for entry in hash_table.entries() {
1499 table_data.write_u32_le(entry.name_1)?;
1500 table_data.write_u32_le(entry.name_2)?;
1501 table_data.write_u16_le(entry.locale)?;
1502 table_data.write_u16_le(entry.platform)?;
1503 table_data.write_u32_le(entry.block_index)?;
1504 }
1505
1506 let key = hash_string("(hash table)", hash_type::FILE_KEY);
1508 self.encrypt_data(&mut table_data, key);
1509
1510 let md5 = self.calculate_md5(&table_data);
1512
1513 writer.write_all(&table_data)?;
1515
1516 Ok(md5)
1517 }
1518
1519 fn write_block_table<W: Write>(
1521 &self,
1522 writer: &mut W,
1523 block_table: &BlockTable,
1524 ) -> Result<[u8; 16]> {
1525 let mut table_data = Vec::new();
1527 for entry in block_table.entries() {
1528 table_data.write_u32_le(entry.file_pos)?;
1529 table_data.write_u32_le(entry.compressed_size)?;
1530 table_data.write_u32_le(entry.file_size)?;
1531 table_data.write_u32_le(entry.flags)?;
1532 }
1533
1534 let key = hash_string("(block table)", hash_type::FILE_KEY);
1536 self.encrypt_data(&mut table_data, key);
1537
1538 let md5 = self.calculate_md5(&table_data);
1540
1541 writer.write_all(&table_data)?;
1543
1544 Ok(md5)
1545 }
1546
1547 fn write_hi_block_table<W: Write>(
1549 &self,
1550 writer: &mut W,
1551 hi_block_table: &HiBlockTable,
1552 ) -> Result<[u8; 16]> {
1553 let mut table_data = Vec::new();
1555 for &entry in hi_block_table.entries() {
1556 table_data.write_u16_le(entry)?;
1557 }
1558
1559 let md5 = self.calculate_md5(&table_data);
1561
1562 writer.write_all(&table_data)?;
1564
1565 Ok(md5)
1566 }
1567
1568 fn write_header<W: Write + Seek>(
1570 &self,
1571 writer: &mut W,
1572 params: &HeaderWriteParams,
1573 ) -> Result<()> {
1574 writer.write_u32_le(crate::signatures::MPQ_ARCHIVE)?;
1576
1577 writer.write_u32_le(self.version.header_size())?;
1579
1580 writer.write_u32_le(params.archive_size.min(u32::MAX as u64) as u32)?;
1582
1583 writer.write_u16_le(self.version as u16)?;
1585
1586 writer.write_u16_le(self.block_size)?;
1588
1589 writer.write_u32_le(params.hash_table_pos as u32)?;
1591 writer.write_u32_le(params.block_table_pos as u32)?;
1592 writer.write_u32_le(params.hash_table_size)?;
1593 writer.write_u32_le(params.block_table_size)?;
1594
1595 match self.version {
1597 FormatVersion::V1 => {
1598 }
1600 FormatVersion::V2 => {
1601 writer.write_u64_le(params.hi_block_table_pos.unwrap_or(0))?;
1603
1604 writer.write_u16_le((params.hash_table_pos >> 32) as u16)?; writer.write_u16_le((params.block_table_pos >> 32) as u16)?; }
1608 FormatVersion::V3 => {
1609 writer.write_u64_le(params.hi_block_table_pos.unwrap_or(0))?; writer.write_u16_le((params.hash_table_pos >> 32) as u16)?; writer.write_u16_le((params.block_table_pos >> 32) as u16)?; writer.write_u64_le(params.archive_size)?; writer.write_u64_le(params.het_table_pos.unwrap_or(0))?; writer.write_u64_le(params.bet_table_pos.unwrap_or(0))?; }
1619 FormatVersion::V4 => {
1620 writer.write_u64_le(params.hi_block_table_pos.unwrap_or(0))?; writer.write_u16_le((params.hash_table_pos >> 32) as u16)?; writer.write_u16_le((params.block_table_pos >> 32) as u16)?; writer.write_u64_le(params.archive_size)?; writer.write_u64_le(params.het_table_pos.unwrap_or(0))?; writer.write_u64_le(params.bet_table_pos.unwrap_or(0))?; if let Some(v4_data) = ¶ms.v4_data {
1632 writer.write_u64_le(v4_data.hash_table_size_64)?;
1633 writer.write_u64_le(v4_data.block_table_size_64)?;
1634 writer.write_u64_le(v4_data.hi_block_table_size_64)?;
1635 writer.write_u64_le(v4_data.het_table_size_64)?;
1636 writer.write_u64_le(v4_data.bet_table_size_64)?;
1637 writer.write_u32_le(v4_data.raw_chunk_size)?;
1638
1639 writer.write_all(&v4_data.md5_block_table)?;
1641 writer.write_all(&v4_data.md5_hash_table)?;
1642 writer.write_all(&v4_data.md5_hi_block_table)?;
1643 writer.write_all(&v4_data.md5_bet_table)?;
1644 writer.write_all(&v4_data.md5_het_table)?;
1645 writer.write_all(&v4_data.md5_mpq_header)?;
1646 } else {
1647 return Err(Error::invalid_format("V4 format requires v4_data"));
1648 }
1649 }
1650 }
1651
1652 Ok(())
1653 }
1654
1655 fn calculate_file_key(&self, filename: &str, file_pos: u64, file_size: u32, flags: u32) -> u32 {
1657 let base_key = hash_string(filename, hash_type::FILE_KEY);
1658
1659 if flags & BlockEntry::FLAG_FIX_KEY != 0 {
1660 (base_key.wrapping_add(file_pos as u32)) ^ file_size
1662 } else {
1663 base_key
1664 }
1665 }
1666
1667 pub fn encrypt_data(&self, data: &mut [u8], key: u32) {
1669 if data.is_empty() || key == 0 {
1670 return;
1671 }
1672
1673 let (chunks, remainder) = data.split_at_mut((data.len() / 4) * 4);
1675
1676 let mut u32_buffer = Vec::with_capacity(chunks.len() / 4);
1678 for chunk in chunks.chunks_exact(4) {
1679 u32_buffer.push(u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
1680 }
1681
1682 encrypt_block(&mut u32_buffer, key);
1683
1684 for (i, &encrypted) in u32_buffer.iter().enumerate() {
1686 let bytes = encrypted.to_le_bytes();
1687 chunks[i * 4..(i + 1) * 4].copy_from_slice(&bytes);
1688 }
1689
1690 if !remainder.is_empty() {
1692 let mut last_dword = [0u8; 4];
1693 last_dword[..remainder.len()].copy_from_slice(remainder);
1694
1695 let mut last_u32 = u32::from_le_bytes(last_dword);
1696 encrypt_block(
1697 std::slice::from_mut(&mut last_u32),
1698 key.wrapping_add((chunks.len() / 4) as u32),
1699 );
1700
1701 let encrypted_bytes = last_u32.to_le_bytes();
1702 remainder.copy_from_slice(&encrypted_bytes[..remainder.len()]);
1703 }
1704 }
1705 fn encrypt_data_u32(&self, data: &mut [u32], key: u32) {
1707 encrypt_block(data, key);
1708 }
1709
1710 fn calculate_md5(&self, data: &[u8]) -> [u8; 16] {
1712 let mut hasher = Md5::new();
1713 hasher.update(data);
1714 hasher.finalize().into()
1715 }
1716
1717 fn finalize_v4_header_md5<W: Write + Seek + Read>(&self, writer: &mut W) -> Result<()> {
1719 writer.seek(SeekFrom::Start(0))?;
1721 let header_size = self.version.header_size() as usize;
1722 let md5_size = 16;
1723 let header_data_size = header_size - md5_size; let mut header_data = vec![0u8; header_data_size];
1726 writer.read_exact(&mut header_data)?;
1727
1728 let header_md5 = self.calculate_md5(&header_data);
1730
1731 writer.seek(SeekFrom::Start(192))?;
1733 writer.write_all(&header_md5)?;
1734
1735 Ok(())
1736 }
1737
1738 fn create_het_table_with_hash_table(
1740 &self,
1741 hash_table: &HashTable,
1742 ) -> Result<(Vec<u8>, HetHeader)> {
1743 let mut file_count = 0u32;
1745
1746 for entry in hash_table.entries() {
1748 if !entry.is_empty() {
1749 file_count += 1;
1750 }
1751 }
1752
1753 let hash_table_entries = (file_count * 2).next_power_of_two();
1755
1756 log::debug!(
1757 "Creating HET table from hash_table: {file_count} files, {hash_table_entries} hash entries"
1758 );
1759
1760 let header = HetHeader {
1762 table_size: 0, max_file_count: file_count,
1764 hash_table_size: hash_table_entries, hash_entry_size: 8, total_index_size: hash_table_entries * Self::calculate_bits_needed(file_count as u64),
1767 index_size_extra: 0,
1768 index_size: Self::calculate_bits_needed(file_count as u64),
1769 block_table_size: 0, };
1771
1772 let index_size = header.index_size;
1774
1775 let mut het_hash_table = vec![0xFFu8; hash_table_entries as usize]; let file_indices_size = (header.total_index_size as usize).div_ceil(8);
1780 let mut file_indices = vec![0u8; file_indices_size]; let invalid_index = (1u64 << index_size) - 1; for i in 0..hash_table_entries {
1785 self.write_bit_entry(&mut file_indices, i as usize, invalid_index, index_size)?;
1786 }
1787
1788 let mut file_index = 0;
1790
1791 let collect_attributes = matches!(
1793 self.attributes_option,
1794 AttributesOption::GenerateCrc32 | AttributesOption::GenerateFull
1795 );
1796
1797 for pending_file in self.pending_files.iter() {
1798 if pending_file.archive_name == "(attributes)" && collect_attributes {
1800 continue;
1801 }
1802
1803 let hash_bits = 8;
1804 let (hash, name_hash1) = het_hash(&pending_file.archive_name, hash_bits);
1805
1806 let start_index = (hash % hash_table_entries as u64) as usize;
1808
1809 let mut current_index = start_index;
1811 loop {
1812 if het_hash_table[current_index] == 0xFF {
1814 het_hash_table[current_index] = name_hash1;
1816
1817 self.write_bit_entry(
1819 &mut file_indices,
1820 current_index,
1821 file_index as u64,
1822 index_size,
1823 )?;
1824
1825 break;
1826 }
1827
1828 current_index = (current_index + 1) % hash_table_entries as usize;
1829 if current_index == start_index {
1830 return Err(Error::invalid_format("HET table full"));
1831 }
1832 }
1833 file_index += 1;
1834 }
1835
1836 if collect_attributes {
1838 let hash_bits = 8;
1839 let (hash, name_hash1) = het_hash("(attributes)", hash_bits);
1840 let start_index = (hash % hash_table_entries as u64) as usize;
1841
1842 let mut current_index = start_index;
1843 loop {
1844 if het_hash_table[current_index] == 0xFF {
1845 het_hash_table[current_index] = name_hash1;
1846 self.write_bit_entry(
1847 &mut file_indices,
1848 current_index,
1849 file_index as u64,
1850 index_size,
1851 )?;
1852 break;
1853 }
1854
1855 current_index = (current_index + 1) % hash_table_entries as usize;
1856 if current_index == start_index {
1857 return Err(Error::invalid_format("HET table full"));
1858 }
1859 }
1860 }
1861
1862 let het_header_size = std::mem::size_of::<HetHeader>();
1864 let data_size = het_header_size as u32 + hash_table_entries + file_indices_size as u32;
1865 let table_size = 12 + data_size; let mut final_header = header;
1869 final_header.table_size = table_size;
1870
1871 let mut result = Vec::with_capacity((12 + data_size) as usize);
1873
1874 result.write_u32_le(0x1A544548)?; result.write_u32_le(1)?; result.write_u32_le(data_size)?; result.write_u32_le(final_header.table_size)?;
1881 result.write_u32_le(final_header.max_file_count)?;
1882 result.write_u32_le(final_header.hash_table_size)?;
1883 result.write_u32_le(final_header.hash_entry_size)?;
1884 result.write_u32_le(final_header.total_index_size)?;
1885 result.write_u32_le(final_header.index_size_extra)?;
1886 result.write_u32_le(final_header.index_size)?;
1887 result.write_u32_le(final_header.block_table_size)?;
1888
1889 result.extend_from_slice(&het_hash_table);
1891 result.extend_from_slice(&file_indices);
1892
1893 log::debug!("HET table created with {file_count} files (including attributes)");
1894
1895 Ok((result, final_header))
1896 }
1897
1898 fn write_bit_entry(
1900 &self,
1901 data: &mut [u8],
1902 index: usize,
1903 value: u64,
1904 bit_size: u32,
1905 ) -> Result<()> {
1906 let bit_offset = index * bit_size as usize;
1907 let byte_offset = bit_offset / 8;
1908 let bit_shift = bit_offset % 8;
1909
1910 let bits_needed = bit_shift + bit_size as usize;
1912 let bytes_needed = bits_needed.div_ceil(8);
1913
1914 if byte_offset + bytes_needed > data.len() {
1915 log::error!(
1916 "Bit entry out of bounds: index={}, bit_size={}, bit_offset={}, byte_offset={}, bytes_needed={}, data.len()={}",
1917 index,
1918 bit_size,
1919 bit_offset,
1920 byte_offset,
1921 bytes_needed,
1922 data.len()
1923 );
1924 return Err(Error::invalid_format("Bit entry out of bounds"));
1925 }
1926
1927 let mut existing = 0u64;
1929 let max_bytes = bytes_needed.min(8);
1930 for i in 0..max_bytes {
1931 if byte_offset + i < data.len() && i * 8 < 64 {
1932 existing |= (data[byte_offset + i] as u64) << (i * 8);
1933 }
1934 }
1935
1936 let value_mask = if bit_size >= 64 {
1938 u64::MAX
1939 } else {
1940 (1u64 << bit_size) - 1
1941 };
1942 let mask = value_mask << bit_shift;
1943 existing &= !mask;
1944
1945 existing |= (value & value_mask) << bit_shift;
1947
1948 for i in 0..max_bytes {
1950 if byte_offset + i < data.len() && i * 8 < 64 {
1951 data[byte_offset + i] = (existing >> (i * 8)) as u8;
1952 }
1953 }
1954
1955 Ok(())
1956 }
1957
1958 fn calculate_bits_needed(max_value: u64) -> u32 {
1960 if max_value == 0 {
1961 1
1962 } else {
1963 (64 - max_value.leading_zeros()).max(1)
1964 }
1965 }
1966
1967 fn write_het_table<W: Write>(
1969 &self,
1970 writer: &mut W,
1971 data: &[u8],
1972 encrypt: bool,
1973 ) -> Result<(u64, [u8; 16])> {
1974 if data.len() < 12 {
1979 return Err(Error::invalid_format("HET table data too small"));
1980 }
1981
1982 let (extended_header, table_data) = data.split_at(12);
1984 let mut processed_data = table_data.to_vec();
1985
1986 log::debug!(
1987 "HET table before processing: extended_header len={}, table_data len={}",
1988 extended_header.len(),
1989 processed_data.len()
1990 );
1991 log::debug!(
1992 "HET table data (first 20 bytes): {:?}",
1993 &processed_data[..processed_data.len().min(20)]
1994 );
1995
1996 if self.compress_tables && matches!(self.version, FormatVersion::V3 | FormatVersion::V4) {
1998 log::debug!("Compressing HET table data: {} -> ", processed_data.len());
1999 let compressed = compress(&processed_data, self.table_compression)?;
2000 log::debug!(
2001 "{} bytes ({}% reduction)",
2002 compressed.len(),
2003 (100 * (processed_data.len() - compressed.len()) / processed_data.len())
2004 );
2005
2006 let mut compressed_with_type = Vec::with_capacity(1 + compressed.len());
2008 compressed_with_type.push(self.table_compression);
2009 compressed_with_type.extend_from_slice(&compressed);
2010 processed_data = compressed_with_type;
2011 }
2012
2013 if encrypt {
2015 let key = hash_string("(hash table)", hash_type::FILE_KEY);
2016 log::debug!("Encrypting HET table with key: 0x{key:08X}");
2017 log::debug!(
2018 "Data size: {} bytes (multiple of 4: {})",
2019 processed_data.len(),
2020 processed_data.len() % 4 == 0
2021 );
2022 log::debug!(
2023 "Data before encryption (last 10 bytes): {:?}",
2024 &processed_data[processed_data.len().saturating_sub(10)..]
2025 );
2026 self.encrypt_data(&mut processed_data, key);
2027 log::debug!(
2028 "Data after encryption (last 10 bytes): {:?}",
2029 &processed_data[processed_data.len().saturating_sub(10)..]
2030 );
2031 }
2032
2033 let mut final_data = Vec::with_capacity(extended_header.len() + processed_data.len());
2035 final_data.extend_from_slice(extended_header);
2036 final_data.extend_from_slice(&processed_data);
2037
2038 let md5 = self.calculate_md5(&final_data);
2040
2041 let written_size = final_data.len() as u64;
2042 writer.write_all(&final_data)?;
2043 Ok((written_size, md5))
2044 }
2045
2046 fn create_bet_table(&self, block_table: &BlockTable) -> Result<(Vec<u8>, BetHeader)> {
2048 let file_count = block_table.entries().len() as u32;
2050
2051 let mut max_file_pos = 0u64;
2053 let mut max_file_size = 0u64;
2054 let mut max_compressed_size = 0u64;
2055 let mut unique_flags = std::collections::HashSet::new();
2056
2057 for i in 0..file_count as usize {
2058 if let Some(entry) = block_table.get(i) {
2059 max_file_pos = max_file_pos.max(entry.file_pos as u64);
2060 max_file_size = max_file_size.max(entry.file_size as u64);
2061 max_compressed_size = max_compressed_size.max(entry.compressed_size as u64);
2062 unique_flags.insert(entry.flags);
2063 }
2064 }
2065
2066 let bit_count_file_pos = Self::calculate_bits_needed(max_file_pos);
2068 let bit_count_file_size = Self::calculate_bits_needed(max_file_size);
2069 let bit_count_cmp_size = Self::calculate_bits_needed(max_compressed_size);
2070 let bit_count_flag_index = if unique_flags.is_empty() {
2071 0
2072 } else {
2073 Self::calculate_bits_needed(unique_flags.len() as u64 - 1)
2074 };
2075 let bit_count_unknown = 0; let bit_index_file_pos = 0;
2079 let bit_index_file_size = bit_index_file_pos + bit_count_file_pos;
2080 let bit_index_cmp_size = bit_index_file_size + bit_count_file_size;
2081 let bit_index_flag_index = bit_index_cmp_size + bit_count_cmp_size;
2082 let bit_index_unknown = bit_index_flag_index + bit_count_flag_index;
2083
2084 let table_entry_size = bit_index_unknown + bit_count_unknown;
2086
2087 let mut flag_array: Vec<u32> = unique_flags.into_iter().collect();
2089 flag_array.sort();
2090 let flag_count = flag_array.len() as u32;
2091
2092 let mut flag_index_map = std::collections::HashMap::new();
2094 for (index, &flags) in flag_array.iter().enumerate() {
2095 flag_index_map.insert(flags, index as u32);
2096 }
2097
2098 let file_table_bits = file_count * table_entry_size;
2100 let file_table_size = file_table_bits.div_ceil(8); let bet_hash_size = 64;
2104 let total_bet_hash_size = file_count * bet_hash_size;
2105 let bet_hash_size_extra = 0;
2106 let bet_hash_array_size = total_bet_hash_size.div_ceil(8);
2107
2108 let header = BetHeader {
2110 table_size: 0, file_count,
2112 unknown_08: 0x10,
2113 table_entry_size,
2114 bit_index_file_pos,
2115 bit_index_file_size,
2116 bit_index_cmp_size,
2117 bit_index_flag_index,
2118 bit_index_unknown,
2119 bit_count_file_pos,
2120 bit_count_file_size,
2121 bit_count_cmp_size,
2122 bit_count_flag_index,
2123 bit_count_unknown,
2124 total_bet_hash_size,
2125 bet_hash_size_extra,
2126 bet_hash_size,
2127 bet_hash_array_size,
2128 flag_count,
2129 };
2130
2131 let mut file_table = vec![0u8; file_table_size as usize];
2133
2134 let mut bet_hashes = Vec::with_capacity(file_count as usize);
2136
2137 for i in 0..file_count as usize {
2139 if let Some(entry) = block_table.get(i) {
2140 let flag_index = flag_index_map.get(&entry.flags).unwrap();
2142
2143 let mut entry_bits = 0u64;
2145 entry_bits |= (entry.file_pos as u64) << bit_index_file_pos;
2146 entry_bits |= (entry.file_size as u64) << bit_index_file_size;
2147 entry_bits |= (entry.compressed_size as u64) << bit_index_cmp_size;
2148 entry_bits |= (*flag_index as u64) << bit_index_flag_index;
2149
2150 self.write_bit_entry(&mut file_table, i, entry_bits, table_entry_size)?;
2152
2153 let filename = if i < self.pending_files.len() {
2156 &self.pending_files[i].archive_name
2157 } else {
2158 "(attributes)"
2160 };
2161 let hash = jenkins_hash(filename);
2162 bet_hashes.push(hash);
2163 }
2164 }
2165
2166 let bet_header_size = std::mem::size_of::<BetHeader>();
2168 let flag_array_size = flag_count * 4;
2169 let data_size =
2170 bet_header_size as u32 + flag_array_size + file_table_size + bet_hash_array_size;
2171 let table_size = 12 + data_size; let mut final_header = header;
2175 final_header.table_size = table_size;
2176
2177 let mut result = Vec::with_capacity((12 + data_size) as usize);
2179
2180 result.write_u32_le(0x1A544542)?; result.write_u32_le(1)?; result.write_u32_le(data_size)?; result.write_u32_le(final_header.table_size)?;
2187 result.write_u32_le(final_header.file_count)?;
2188 result.write_u32_le(final_header.unknown_08)?;
2189 result.write_u32_le(final_header.table_entry_size)?;
2190 result.write_u32_le(final_header.bit_index_file_pos)?;
2191 result.write_u32_le(final_header.bit_index_file_size)?;
2192 result.write_u32_le(final_header.bit_index_cmp_size)?;
2193 result.write_u32_le(final_header.bit_index_flag_index)?;
2194 result.write_u32_le(final_header.bit_index_unknown)?;
2195 result.write_u32_le(final_header.bit_count_file_pos)?;
2196 result.write_u32_le(final_header.bit_count_file_size)?;
2197 result.write_u32_le(final_header.bit_count_cmp_size)?;
2198 result.write_u32_le(final_header.bit_count_flag_index)?;
2199 result.write_u32_le(final_header.bit_count_unknown)?;
2200 result.write_u32_le(final_header.total_bet_hash_size)?;
2201 result.write_u32_le(final_header.bet_hash_size_extra)?;
2202 result.write_u32_le(final_header.bet_hash_size)?;
2203 result.write_u32_le(final_header.bet_hash_array_size)?;
2204 result.write_u32_le(final_header.flag_count)?;
2205
2206 for &flags in &flag_array {
2208 result.write_u32_le(flags)?;
2209 }
2210
2211 result.extend_from_slice(&file_table);
2213
2214 let mut hash_bytes = vec![0u8; bet_hash_array_size as usize];
2216 for (i, &hash) in bet_hashes.iter().enumerate() {
2217 self.write_bit_entry(&mut hash_bytes, i, hash, bet_hash_size)?;
2218 }
2219 result.extend_from_slice(&hash_bytes);
2220
2221 Ok((result, final_header))
2222 }
2223
2224 fn write_bet_table<W: Write>(
2226 &self,
2227 writer: &mut W,
2228 data: &[u8],
2229 encrypt: bool,
2230 ) -> Result<(u64, [u8; 16])> {
2231 if data.len() < 12 {
2236 return Err(Error::invalid_format("BET table data too small"));
2237 }
2238
2239 let (extended_header, table_data) = data.split_at(12);
2241 let mut processed_data = table_data.to_vec();
2242
2243 if self.compress_tables && matches!(self.version, FormatVersion::V3 | FormatVersion::V4) {
2245 log::debug!("Compressing BET table data: {} -> ", processed_data.len());
2246 let compressed = compress(&processed_data, self.table_compression)?;
2247 log::debug!(
2248 "{} bytes ({}% reduction)",
2249 compressed.len(),
2250 (100 * (processed_data.len() - compressed.len()) / processed_data.len())
2251 );
2252
2253 let mut compressed_with_type = Vec::with_capacity(1 + compressed.len());
2255 compressed_with_type.push(self.table_compression);
2256 compressed_with_type.extend_from_slice(&compressed);
2257 processed_data = compressed_with_type;
2258 }
2259
2260 if encrypt {
2262 let key = hash_string("(block table)", hash_type::FILE_KEY);
2263 self.encrypt_data(&mut processed_data, key);
2264 }
2265
2266 let mut final_data = Vec::with_capacity(extended_header.len() + processed_data.len());
2268 final_data.extend_from_slice(extended_header);
2269 final_data.extend_from_slice(&processed_data);
2270
2271 let md5 = self.calculate_md5(&final_data);
2273
2274 let written_size = final_data.len() as u64;
2275 writer.write_all(&final_data)?;
2276 Ok((written_size, md5))
2277 }
2278}
2279
2280impl Default for ArchiveBuilder {
2281 fn default() -> Self {
2282 Self::new()
2283 }
2284}