1use {
7 anyhow::Result,
8 crc32fast::Hasher,
9 fatfs::{Dir, FileSystem, FormatVolumeOptions, FsOptions, ReadWriteSeek},
10 flate2::{bufread::ZlibEncoder, read::ZlibDecoder, Compression},
11 fscommon::BufStream,
12 gpt::mbr::{PartRecord, ProtectiveMBR},
13 std::{
14 fs::File,
15 io::{BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write},
16 path::Path,
17 },
18};
19
20mod blkx;
21mod koly;
22mod xml;
23
24pub use crate::{blkx::*, koly::*, xml::*};
25
26pub struct DmgReader<R: Read + Seek> {
27 koly: KolyTrailer,
28 xml: Plist,
29 r: R,
30}
31
32impl DmgReader<BufReader<File>> {
33 pub fn open(path: &Path) -> Result<Self> {
34 let r = BufReader::new(File::open(path)?);
35 Self::new(r)
36 }
37}
38
39impl<R: Read + Seek> DmgReader<R> {
40 pub fn new(mut r: R) -> Result<Self> {
41 let koly = KolyTrailer::read_from(&mut r)?;
42 r.seek(SeekFrom::Start(koly.plist_offset))?;
43 let mut xml = Vec::with_capacity(koly.plist_length as usize);
44 (&mut r).take(koly.plist_length).read_to_end(&mut xml)?;
45 let xml: Plist = plist::from_reader_xml(&xml[..])?;
46 Ok(Self { koly, xml, r })
47 }
48
49 pub fn koly(&self) -> &KolyTrailer {
50 &self.koly
51 }
52
53 pub fn plist(&self) -> &Plist {
54 &self.xml
55 }
56
57 pub fn sector(&mut self, chunk: &BlkxChunk) -> Result<impl Read + '_> {
58 self.r.seek(SeekFrom::Start(chunk.compressed_offset))?;
59 let compressed_chunk = (&mut self.r).take(chunk.compressed_length);
60 match chunk.ty().expect("unknown chunk type") {
61 ChunkType::Ignore | ChunkType::Zero | ChunkType::Comment => {
62 Ok(Box::new(std::io::repeat(0).take(chunk.compressed_length)) as Box<dyn Read>)
63 }
64 ChunkType::Raw => Ok(Box::new(compressed_chunk)),
65 ChunkType::Zlib => Ok(Box::new(ZlibDecoder::new(compressed_chunk))),
66 ChunkType::Adc | ChunkType::Bzlib | ChunkType::Lzfse => unimplemented!(),
67 ChunkType::Term => Ok(Box::new(std::io::empty())),
68 }
69 }
70
71 pub fn data_checksum(&mut self) -> Result<u32> {
72 self.r.seek(SeekFrom::Start(self.koly.data_fork_offset))?;
73 let mut data_fork = Vec::with_capacity(self.koly.data_fork_length as usize);
74 (&mut self.r)
75 .take(self.koly.data_fork_length)
76 .read_to_end(&mut data_fork)?;
77 Ok(crc32fast::hash(&data_fork))
78 }
79
80 pub fn partition_table(&self, i: usize) -> Result<BlkxTable> {
81 self.plist().partitions()[i].table()
82 }
83
84 pub fn partition_name(&self, i: usize) -> &str {
85 &self.plist().partitions()[i].name
86 }
87
88 pub fn partition_data(&mut self, i: usize) -> Result<Vec<u8>> {
89 let table = self.plist().partitions()[i].table()?;
90 let mut partition = vec![];
91 for chunk in &table.chunks {
92 std::io::copy(&mut self.sector(chunk)?, &mut partition)?;
93 }
94 Ok(partition)
95 }
96}
97
98pub struct DmgWriter<W: Write + Seek> {
99 xml: Plist,
100 w: W,
101 data_hasher: Hasher,
102 main_hasher: Hasher,
103 sector_number: u64,
104 compressed_offset: u64,
105}
106
107impl DmgWriter<BufWriter<File>> {
108 pub fn create(path: &Path) -> Result<Self> {
109 let w = BufWriter::new(File::create(path)?);
110 Ok(Self::new(w))
111 }
112}
113
114impl<W: Write + Seek> DmgWriter<W> {
115 pub fn new(w: W) -> Self {
116 Self {
117 xml: Default::default(),
118 w,
119 data_hasher: Hasher::new(),
120 main_hasher: Hasher::new(),
121 sector_number: 0,
122 compressed_offset: 0,
123 }
124 }
125
126 pub fn create_fat32(mut self, fat32: &[u8]) -> Result<()> {
127 anyhow::ensure!(fat32.len() % 512 == 0);
128 let sector_count = fat32.len() as u64 / 512;
129 let mut mbr = ProtectiveMBR::new();
130 let mut partition = PartRecord::new_protective(Some(sector_count.try_into()?));
131 partition.os_type = 11;
132 mbr.set_partition(0, partition);
133 let mbr = mbr.to_bytes().to_vec();
134 self.add_partition("Master Boot Record (MBR : 0)", &mbr)?;
135 self.add_partition("FAT32 (FAT32 : 1)", fat32)?;
136 self.finish()?;
137 Ok(())
138 }
139
140 pub fn add_partition(&mut self, name: &str, bytes: &[u8]) -> Result<()> {
141 anyhow::ensure!(bytes.len() % 512 == 0);
142 let id = self.xml.partitions().len() as u32;
143 let name = name.to_string();
144 let mut table = BlkxTable::new(id, self.sector_number, crc32fast::hash(bytes));
145 for chunk in bytes.chunks(2048 * 512) {
146 let mut encoder = ZlibEncoder::new(chunk, Compression::best());
147 let mut compressed = vec![];
148 encoder.read_to_end(&mut compressed)?;
149 let compressed_length = compressed.len() as u64;
150 let sector_count = chunk.len() as u64 / 512;
151 self.w.write_all(&compressed)?;
152 self.data_hasher.update(&compressed);
153 table.add_chunk(BlkxChunk::new(
154 ChunkType::Zlib,
155 self.sector_number,
156 sector_count,
157 self.compressed_offset,
158 compressed_length,
159 ));
160 self.sector_number += sector_count;
161 self.compressed_offset += compressed_length;
162 }
163 table.add_chunk(BlkxChunk::term(self.sector_number, self.compressed_offset));
164 self.main_hasher.update(&table.checksum.data[..4]);
165 self.xml
166 .add_partition(Partition::new(id as i32 - 1, name, table));
167 Ok(())
168 }
169
170 pub fn finish(mut self) -> Result<()> {
171 let mut xml = vec![];
172 plist::to_writer_xml(&mut xml, &self.xml)?;
173 let pos = self.w.stream_position()?;
174 let data_digest = self.data_hasher.finalize();
175 let main_digest = self.main_hasher.finalize();
176 let koly = KolyTrailer::new(
177 pos,
178 self.sector_number,
179 pos,
180 xml.len() as _,
181 data_digest,
182 main_digest,
183 );
184 self.w.write_all(&xml)?;
185 koly.write_to(&mut self.w)?;
186 Ok(())
187 }
188}
189
190fn symlink(target: &str) -> Result<Vec<u8>> {
192 let xsym = format!(
193 "XSym\n{:04}\n{:x}\n{}\n",
194 target.as_bytes().len(),
195 md5::compute(target.as_bytes()),
196 target,
197 );
198 let mut xsym = xsym.into_bytes();
199 anyhow::ensure!(xsym.len() <= 1067);
200 xsym.resize(1067, b' ');
201 Ok(xsym)
202}
203
204fn add_dir<T: ReadWriteSeek>(src: &Path, dest: &Dir<'_, T>) -> Result<()> {
205 for entry in std::fs::read_dir(src)? {
206 let entry = entry?;
207 let file_name = entry.file_name();
208 let file_name = file_name.to_str().unwrap();
209 let source = src.join(file_name);
210 let file_type = entry.file_type()?;
211 if file_type.is_dir() {
212 let d = dest.create_dir(file_name)?;
213 add_dir(&source, &d)?;
214 } else if file_type.is_file() {
215 let mut f = dest.create_file(file_name)?;
216 std::io::copy(&mut File::open(source)?, &mut f)?;
217 } else if file_type.is_symlink() {
218 let target = std::fs::read_link(&source)?;
219 let xsym = symlink(target.to_str().unwrap())?;
220 let mut f = dest.create_file(file_name)?;
221 std::io::copy(&mut &xsym[..], &mut f)?;
222 }
223 }
224 Ok(())
225}
226
227pub fn create_dmg(dir: &Path, dmg: &Path, volume_label: &str, total_sectors: u32) -> Result<()> {
228 let mut fat32 = vec![0; total_sectors as usize * 512];
229 {
230 let mut volume_label_bytes = [0; 11];
231 let end = std::cmp::min(volume_label_bytes.len(), volume_label.len());
232 volume_label_bytes[..end].copy_from_slice(&volume_label.as_bytes()[..end]);
233 let volume_options = FormatVolumeOptions::new()
234 .volume_label(volume_label_bytes)
235 .bytes_per_sector(512)
236 .total_sectors(total_sectors);
237 let mut disk = BufStream::new(Cursor::new(&mut fat32));
238 fatfs::format_volume(&mut disk, volume_options)?;
239 let fs = FileSystem::new(disk, FsOptions::new())?;
240 let file_name = dir.file_name().unwrap().to_str().unwrap();
241 let dest = fs.root_dir().create_dir(file_name)?;
242 add_dir(dir, &dest)?;
243 }
244 DmgWriter::create(dmg)?.create_fat32(&fat32)
245}
246
247#[cfg(test)]
248mod tests {
249 use {super::*, gpt::disk::LogicalBlockSize};
250
251 static DMG: &[u8] = include_bytes!("../assets/example.dmg");
252
253 fn print_dmg<R: Read + Seek>(dmg: &DmgReader<R>) -> Result<()> {
254 println!("{:?}", dmg.koly());
255 println!("{:?}", dmg.plist());
256 for partition in dmg.plist().partitions() {
257 let table = partition.table()?;
258 println!("{table:?}");
259 println!("table checksum 0x{:x}", u32::from(table.checksum));
260 for (i, chunk) in table.chunks.iter().enumerate() {
261 println!("{i} {chunk:?}");
262 }
263 }
264 Ok(())
265 }
266
267 #[test]
268 fn read_koly_trailer() -> Result<()> {
269 let koly = KolyTrailer::read_from(&mut Cursor::new(DMG))?;
270 let mut bytes = [0; 512];
272 koly.write_to(&mut &mut bytes[..])?;
273 let koly2 = KolyTrailer::read_from(&mut Cursor::new(&bytes))?;
274 assert_eq!(koly, koly2);
275 Ok(())
276 }
277
278 #[test]
279 fn only_read_dmg() -> Result<()> {
280 let mut dmg = DmgReader::new(Cursor::new(DMG))?;
281 print_dmg(&dmg)?;
282 assert_eq!(
283 UdifChecksum::new(dmg.data_checksum()?),
284 dmg.koly().data_fork_digest
285 );
286 let mut buffer = vec![];
287 let mut dmg2 = DmgWriter::new(Cursor::new(&mut buffer));
288 for i in 0..dmg.plist().partitions().len() {
289 let data = dmg.partition_data(i)?;
290 let name = dmg.partition_name(i);
291 dmg2.add_partition(name, &data)?;
292 }
293 dmg2.finish()?;
294 let mut dmg2 = DmgReader::new(Cursor::new(buffer))?;
295 print_dmg(&dmg2)?;
296 assert_eq!(
297 UdifChecksum::new(dmg2.data_checksum()?),
298 dmg2.koly().data_fork_digest
299 );
300 for i in 0..dmg.plist().partitions().len() {
301 let table = dmg.partition_table(i)?;
302 let data = dmg.partition_data(i)?;
303 let expected = u32::from(table.checksum);
304 let calculated = crc32fast::hash(&data);
305 assert_eq!(expected, calculated);
306 }
307 assert_eq!(dmg.koly().main_digest, dmg2.koly().main_digest);
308 println!("data crc32 0x{:x}", u32::from(dmg.koly().data_fork_digest));
309 println!("main crc32 0x{:x}", u32::from(dmg.koly().main_digest));
310 Ok(())
311 }
312
313 #[test]
314 fn read_dmg_partition_mbr() -> Result<()> {
315 let mut dmg = DmgReader::new(Cursor::new(DMG))?;
316 let mbr = dmg.partition_data(0)?;
317 println!("{mbr:?}");
318 let mbr = ProtectiveMBR::from_bytes(&mbr, LogicalBlockSize::Lb512)?;
319 println!("{mbr:?}");
320 Ok(())
321 }
322
323 #[test]
324 fn read_dmg_partition_fat32() -> Result<()> {
325 let mut dmg = DmgReader::new(Cursor::new(DMG))?;
326 let fat32 = dmg.partition_data(1)?;
327 let fs = FileSystem::new(Cursor::new(fat32), FsOptions::new())?;
328 println!("volume: {}", fs.volume_label());
329 for entry in fs.root_dir().iter() {
330 let entry = entry?;
331 println!("{}", entry.file_name());
332 }
333 Ok(())
334 }
335
336 #[test]
337 fn checksum() -> Result<()> {
338 let mut dmg = DmgReader::new(Cursor::new(DMG))?;
339 assert_eq!(
340 UdifChecksum::new(dmg.data_checksum()?),
341 dmg.koly().data_fork_digest
342 );
343 for i in 0..dmg.plist().partitions().len() {
344 let table = dmg.partition_table(i)?;
345 let data = dmg.partition_data(i)?;
346 let expected = u32::from(table.checksum);
347 let calculated = crc32fast::hash(&data);
348 assert_eq!(expected, calculated);
349 }
350 Ok(())
351 }
352}