1use crate::extensions;
2use std::fs::File;
3use std::io::{BufReader, Read};
4use std::path::Path;
5use std::{cmp, io, iter};
6use thiserror::Error;
7use zip::ZipArchive;
8use zip::result::ZipError;
9
10#[derive(Debug, Error)]
11pub enum ArchiveError {
12    #[error("I/O error reading archive file '{path}': {source}")]
13    Io {
14        path: String,
15        #[source]
16        source: io::Error,
17    },
18    #[error("Error reading .zip archive '{path}': {source}")]
19    Zip {
20        path: String,
21        #[source]
22        source: ZipError,
23    },
24    #[error("Error reading .7z archive '{path}': {source}")]
25    SevenZ {
26        path: String,
27        #[source]
28        source: sevenz_rust::Error,
29    },
30    #[error("File '{file_name}' not found in archive '{path}'")]
31    FileNotFound { path: String, file_name: String },
32    #[error("No files with supported extensions found in archive '{path}'")]
33    NoSupportedFiles { path: String },
34}
35
36impl ArchiveError {
37    fn io(path: &Path, source: io::Error) -> Self {
38        Self::Io { path: path.display().to_string(), source }
39    }
40
41    fn zip(path: &Path, source: ZipError) -> Self {
42        Self::Zip { path: path.display().to_string(), source }
43    }
44
45    fn sevenz(path: &Path, source: sevenz_rust::Error) -> Self {
46        Self::SevenZ { path: path.display().to_string(), source }
47    }
48
49    fn file_not_found(path: &Path, file_name: &str) -> Self {
50        Self::FileNotFound { path: path.display().to_string(), file_name: file_name.into() }
51    }
52}
53
54#[derive(Debug, Clone)]
55pub struct ZipEntryMetadata {
56    pub file_name: String,
57    pub extension: String,
58    pub size: u64,
59}
60
61fn extension_matches(file_name: &str, target_extension: &str) -> bool {
62    extensions::from_path(file_name).is_some_and(|file_ext| file_ext.as_str() == target_extension)
63}
64
65#[derive(Debug, Clone, Copy)]
66pub struct ArchiveEntry<'a> {
67    pub file_name: &'a str,
68    pub size: u64,
69}
70
71/// List all files in a .zip archive.
72///
73/// The given callback will be called for every file entry unless there is an error.
74///
75/// # Errors
76///
77/// Propagates any I/O or decoding errors.
78pub fn list_files_zip(
79    zip_path: &Path,
80    mut callback: impl FnMut(ArchiveEntry<'_>),
81) -> Result<(), ArchiveError> {
82    let io_err_fn = |source| ArchiveError::io(zip_path, source);
83    let zip_err_fn = |source| ArchiveError::zip(zip_path, source);
84
85    let file = File::open(zip_path).map_err(io_err_fn)?;
86    let reader = BufReader::new(file);
87    let mut archive = ZipArchive::new(reader).map_err(zip_err_fn)?;
88
89    for i in 0..archive.len() {
90        let entry = archive.by_index(i).map_err(zip_err_fn)?;
91        callback(ArchiveEntry { file_name: entry.name(), size: entry.size() });
92    }
93
94    Ok(())
95}
96
97/// List all files in a .7z archive.
98///
99/// The given callback will be called for every file entry unless there is an error.
100///
101/// # Errors
102///
103/// Propagates any I/O or decoding errors.
104pub fn list_files_7z(
105    sevenz_path: &Path,
106    mut callback: impl FnMut(ArchiveEntry<'_>),
107) -> Result<(), ArchiveError> {
108    let io_err_fn = |source| ArchiveError::io(sevenz_path, source);
109    let sevenz_err_fn = |source| ArchiveError::sevenz(sevenz_path, source);
110
111    let file = File::open(sevenz_path).map_err(io_err_fn)?;
112    let file_len = file.metadata().map_err(io_err_fn)?.len();
113    let mut reader = BufReader::new(file);
114    let archive = sevenz_rust::Archive::read(&mut reader, file_len, &[]).map_err(sevenz_err_fn)?;
115
116    for entry in &archive.files {
117        if !entry.has_stream {
118            // Not a readable file (e.g. is a directory)
119            continue;
120        }
121
122        callback(ArchiveEntry { file_name: entry.name.as_str(), size: entry.size });
123    }
124
125    Ok(())
126}
127
128/// Read a file from within a .zip archive.
129///
130/// # Errors
131///
132/// Propagates any I/O or decoding errors.
133///
134/// Will return an error if the archive does not contain the specified file.
135pub fn read_file_zip(
136    zip_path: &Path,
137    file_name: &str,
138    max_len: usize,
139) -> Result<Vec<u8>, ArchiveError> {
140    let io_err_fn = |source| ArchiveError::io(zip_path, source);
141    let zip_err_fn = |source| ArchiveError::zip(zip_path, source);
142
143    let file = File::open(zip_path).map_err(io_err_fn)?;
144    let reader = BufReader::new(file);
145    let mut archive = ZipArchive::new(reader).map_err(zip_err_fn)?;
146
147    let mut entry = archive.by_name(file_name).map_err(zip_err_fn)?;
148    let len = cmp::min(entry.size() as usize, max_len);
149
150    let mut buffer = vec![0; len];
151    entry.read_exact(&mut buffer).map_err(io_err_fn)?;
152
153    Ok(buffer)
154}
155
156/// Read a file from within a .7z archive.
157///
158/// # Errors
159///
160/// Propagates any I/O or decoding errors.
161///
162/// Will return an error if the archive does not contain the specified file.
163pub fn read_file_7z(
164    sevenz_path: &Path,
165    file_name: &str,
166    max_len: usize,
167) -> Result<Vec<u8>, ArchiveError> {
168    let io_err_fn = |source| ArchiveError::io(sevenz_path, source);
169    let sevenz_err_fn = |source| ArchiveError::sevenz(sevenz_path, source);
170
171    let file = File::open(sevenz_path).map_err(io_err_fn)?;
172    let file_len = file.metadata().map_err(io_err_fn)?.len();
173    let mut reader = BufReader::new(file);
174    let archive = sevenz_rust::Archive::read(&mut reader, file_len, &[]).map_err(sevenz_err_fn)?;
175
176    for folder_idx in 0..archive.folders.len() {
177        let decoder = sevenz_rust::BlockDecoder::new(folder_idx, &archive, &[], &mut reader);
178
179        if !decoder.entries().iter().any(|entry| entry.name.as_str() == file_name) {
180            continue;
181        }
182
183        let mut buffer = Vec::new();
184        let mut found = false;
185        decoder
186            .for_each_entries(&mut |entry, reader| {
187                if entry.name.as_str() == file_name {
188                    let len = cmp::min(entry.size as usize, max_len);
189                    buffer.extend(iter::repeat_n(0, len));
190                    reader.read_exact(&mut buffer[..len])?;
191                    found = true;
192                    Ok(false)
193                } else {
194                    io::copy(reader, &mut io::sink())?;
195                    Ok(true)
196                }
197            })
198            .map_err(sevenz_err_fn)?;
199
200        if found {
201            return Ok(buffer);
202        }
203    }
204
205    Err(ArchiveError::file_not_found(sevenz_path, file_name))
206}
207
208/// Returns metadata of the first file in the .zip archive that has a supported extension, or
209/// None if there are no files with a supported extension.
210///
211/// Will also return None if the archive contains any .cue files, under the assumption that the
212/// archive contains a CD-ROM image.
213///
214/// # Errors
215///
216/// Propagates any I/O or ZIP errors.
217pub fn first_supported_file_in_zip(
218    zip_path: &Path,
219    supported_extensions: &[&str],
220) -> Result<Option<ZipEntryMetadata>, ArchiveError> {
221    let io_err_fn = |source| ArchiveError::io(zip_path, source);
222    let zip_err_fn = |source| ArchiveError::zip(zip_path, source);
223
224    let file = File::open(zip_path).map_err(io_err_fn)?;
225    let reader = BufReader::new(file);
226    let mut archive = ZipArchive::new(reader).map_err(zip_err_fn)?;
227
228    if archive.file_names().any(|file_name| extension_matches(file_name, "cue")) {
229        // Archive contains a .cue file; assume it's a CD-ROM image
230        return Ok(None);
231    }
232
233    let mut first_file_name_with_ext: Option<(String, String)> = None;
234    for file_name in archive.file_names() {
235        let Some(extension) = extensions::from_path(file_name) else {
236            continue;
237        };
238
239        if supported_extensions.contains(&extension.as_str()) {
240            first_file_name_with_ext = Some((file_name.into(), extension));
241            break;
242        }
243    }
244
245    let Some((file_name, extension)) = first_file_name_with_ext else {
246        return Ok(None);
247    };
248
249    let file = archive.by_name(&file_name).map_err(zip_err_fn)?;
250    let size = file.size();
251
252    Ok(Some(ZipEntryMetadata { file_name, extension, size }))
253}