1use std::collections::BTreeMap;
4use std::io::Read;
5use std::ops::ControlFlow;
6use std::path::Path;
7
8use md5::{Digest, Md5};
9use sha1::Sha1;
10
11use super::{
12 EwfWriter, SequentialOptions, SequentialWriter, WriteFormat, WriteOptions, WriteResult,
13 hex_string,
14};
15use crate::{
16 EwfError, MediaType, Result, SingleFileEntry, SingleFileEntryType, SingleFileExtent,
17 SingleFilesInfo,
18};
19
20#[derive(Debug, Clone, Default, PartialEq, Eq)]
25pub struct LogicalEntryMetadata {
26 pub name: String,
28 pub creation_time: Option<i64>,
30 pub modification_time: Option<i64>,
32 pub access_time: Option<i64>,
34 pub entry_modification_time: Option<i64>,
36 pub deletion_time: Option<i64>,
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42#[non_exhaustive]
43pub struct LogicalWriteProgress {
44 pub identifier: u64,
46 pub bytes_written: u64,
48 pub bytes_total: u64,
50}
51
52pub struct LogicalWriter {
58 writer: Backend,
59 paths: BTreeMap<u64, Vec<usize>>,
60 next_identifier: u64,
61 poisoned: bool,
62}
63
64enum Backend {
65 General(Box<EwfWriter>),
66 Sequential(Box<SequentialWriter>),
67}
68
69impl Backend {
70 fn position(&self) -> u64 {
71 match self {
72 Self::General(w) => w.position(),
73 Self::Sequential(w) => w.position(),
74 }
75 }
76 fn options(&mut self) -> &mut WriteOptions {
77 match self {
78 Self::General(w) => &mut w.options,
79 Self::Sequential(w) => &mut w.options,
80 }
81 }
82 fn write_all(&mut self, data: &[u8]) -> Result<()> {
83 match self {
84 Self::General(w) => w.write_all(data),
85 Self::Sequential(w) => w.write_all(data),
86 }
87 }
88 fn finish(self) -> Result<WriteResult> {
89 match self {
90 Self::General(w) => w.finish(),
91 Self::Sequential(w) => w.finish(),
92 }
93 }
94}
95
96impl LogicalWriter {
97 pub fn create(path: impl AsRef<Path>, mut options: WriteOptions) -> Result<Self> {
101 Self::catalog(&mut options)?;
102 Ok(Self::from_backend(Backend::General(Box::new(
103 EwfWriter::create(path, options)?,
104 ))))
105 }
106
107 pub fn create_sequential(
111 path: impl AsRef<Path>,
112 mut options: SequentialOptions,
113 ) -> Result<Self> {
114 if options.write.format != WriteFormat::Ewf2Logical {
115 return Err(EwfError::Unsupported(
116 "sequential logical builder requires EWF2 logical format".into(),
117 ));
118 }
119 Self::catalog(&mut options.write)?;
120 Ok(Self::from_backend(Backend::Sequential(Box::new(
121 SequentialWriter::create(path, options)?,
122 ))))
123 }
124
125 fn from_backend(writer: Backend) -> Self {
126 Self {
127 writer,
128 paths: BTreeMap::from([(1, vec![])]),
129 next_identifier: 2,
130 poisoned: false,
131 }
132 }
133
134 fn catalog(options: &mut WriteOptions) -> Result<()> {
135 if !matches!(
136 options.format,
137 WriteFormat::Ewf1Logical | WriteFormat::Ewf2Logical
138 ) || options.single_files.is_some()
139 || !options.ewf2_single_files_tables.is_empty()
140 {
141 return Err(EwfError::Unsupported(
142 "logical builder requires a logical format and no prebuilt catalog".into(),
143 ));
144 }
145 options.media_profile.media_type = Some(MediaType::SingleFiles);
146 options.single_files = Some(SingleFilesInfo {
147 root: SingleFileEntry {
148 identifier: Some(1),
149 name: Some(String::new()),
150 file_entry_type: Some(SingleFileEntryType::Directory),
151 permission_group_index: Some(-1),
152 ..SingleFileEntry::default()
153 },
154 ..SingleFilesInfo::default()
155 });
156 Ok(())
157 }
158
159 pub fn add_directory(&mut self, parent: u64, metadata: LogicalEntryMetadata) -> Result<u64> {
162 let entry = self.prepare(parent, metadata, SingleFileEntryType::Directory)?;
163 Ok(self.insert(parent, entry))
164 }
165
166 pub fn add_file(
170 &mut self,
171 parent: u64,
172 metadata: LogicalEntryMetadata,
173 size: u64,
174 input: &mut impl Read,
175 ) -> Result<u64> {
176 self.add_file_with_progress(parent, metadata, size, input, |_| ControlFlow::Continue(()))
177 }
178
179 pub fn add_file_with_progress(
182 &mut self,
183 parent: u64,
184 metadata: LogicalEntryMetadata,
185 size: u64,
186 input: &mut impl Read,
187 mut progress: impl FnMut(LogicalWriteProgress) -> ControlFlow<()>,
188 ) -> Result<u64> {
189 let mut entry = self.prepare(parent, metadata, SingleFileEntryType::File)?;
190 let offset = self.writer.position();
191 let end = offset
192 .checked_add(size)
193 .filter(|end| i64::try_from(*end).is_ok())
194 .ok_or_else(|| {
195 EwfError::Malformed("logical file address exceeds EWF catalog range".into())
196 })?;
197 self.poisoned = true;
198 let mut written = 0;
199 let mut buffer = vec![0; 1024 * 1024];
200 let mut md5 = Md5::new();
201 let mut sha1 = Sha1::new();
202 loop {
203 if progress(LogicalWriteProgress {
204 identifier: self.next_identifier,
205 bytes_written: written,
206 bytes_total: size,
207 })
208 .is_break()
209 {
210 return Err(EwfError::Aborted);
211 }
212 if written == size {
213 break;
214 }
215 let length = (size - written).min(buffer.len() as u64) as usize;
216 input.read_exact(&mut buffer[..length])?;
217 self.writer.write_all(&buffer[..length])?;
218 md5.update(&buffer[..length]);
219 sha1.update(&buffer[..length]);
220 written += length as u64;
221 }
222 entry.size = Some(size);
223 entry.logical_offset = Some(offset as i64);
224 if size != 0 {
225 entry.extents.push(SingleFileExtent {
226 data_offset: offset,
227 data_size: size,
228 sparse: false,
229 });
230 }
231 entry.md5 = Some(hex_string(&md5.finalize()));
232 entry.sha1 = Some(hex_string(&sha1.finalize()));
233 self.writer
234 .options()
235 .single_files
236 .as_mut()
237 .expect("builder catalog")
238 .data_size = end;
239 let identifier = self.insert(parent, entry);
240 self.poisoned = false;
241 Ok(identifier)
242 }
243
244 pub fn finish(self) -> Result<WriteResult> {
246 self.ensure_healthy()?;
247 self.writer.finish()
248 }
249
250 fn ensure_healthy(&self) -> Result<()> {
251 if self.poisoned {
252 return Err(EwfError::Unsupported(
253 "logical acquisition failed; discard this writer".into(),
254 ));
255 }
256 Ok(())
257 }
258
259 fn prepare(
260 &self,
261 parent: u64,
262 metadata: LogicalEntryMetadata,
263 kind: SingleFileEntryType,
264 ) -> Result<SingleFileEntry> {
265 self.ensure_healthy()?;
266 if metadata.name.is_empty() || metadata.name.contains(['\0', '\t', '\r', '\n']) {
267 return Err(EwfError::Malformed(
268 "logical entry name is empty or contains catalog delimiters".into(),
269 ));
270 }
271 let path = self
272 .paths
273 .get(&parent)
274 .ok_or_else(|| EwfError::Malformed("logical parent does not exist".into()))?;
275 if path.len() >= 128 {
276 return Err(EwfError::Unsupported(
277 "logical directory nesting exceeds 128".into(),
278 ));
279 }
280 self.next_identifier
281 .checked_add(1)
282 .ok_or_else(|| EwfError::Malformed("logical identifier overflow".into()))?;
283 Ok(SingleFileEntry {
284 identifier: Some(self.next_identifier),
285 name: Some(metadata.name),
286 file_entry_type: Some(kind),
287 permission_group_index: Some(-1),
288 creation_time: metadata.creation_time,
289 modification_time: metadata.modification_time,
290 access_time: metadata.access_time,
291 entry_modification_time: metadata.entry_modification_time,
292 deletion_time: metadata.deletion_time,
293 ..SingleFileEntry::default()
294 })
295 }
296
297 fn insert(&mut self, parent: u64, entry: SingleFileEntry) -> u64 {
298 let mut path = self.paths[&parent].clone();
299 let mut node = &mut self
300 .writer
301 .options()
302 .single_files
303 .as_mut()
304 .expect("builder catalog")
305 .root;
306 for &index in &path {
307 node = &mut node.children[index];
308 }
309 path.push(node.children.len());
310 if entry.file_entry_type == Some(SingleFileEntryType::Directory) {
311 self.paths.insert(self.next_identifier, path);
312 }
313 node.children.push(entry);
314 let identifier = self.next_identifier;
315 self.next_identifier += 1;
316 identifier
317 }
318}