Skip to main content

muy_zipido/
lib.rs

1pub mod circular_buffer;
2pub mod progress_bar;
3
4use circular_buffer::CircularBuffer;
5use flate2::read::DeflateDecoder;
6use progress_bar::ProgressBar;
7use std::error::Error;
8use std::fmt;
9use std::io::Read;
10
11#[derive(Debug)]
12pub enum ZipError {
13    Http(reqwest::Error),
14    UnexpectedEof,
15    InvalidSignature(String),
16    Io(std::io::Error),
17    Decompression(String),
18}
19
20impl fmt::Display for ZipError {
21    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
22        match self {
23            ZipError::Http(e) => write!(f, "HTTP error: {}", e),
24            ZipError::UnexpectedEof => write!(f, "Unexpected end of stream"),
25            ZipError::InvalidSignature(sig) => write!(f, "Invalid signature: {}", sig),
26            ZipError::Io(e) => write!(f, "IO error: {}", e),
27            ZipError::Decompression(e) => write!(f, "Decompression error: {}", e),
28        }
29    }
30}
31
32impl Error for ZipError {}
33
34impl From<reqwest::Error> for ZipError {
35    fn from(e: reqwest::Error) -> Self {
36        ZipError::Http(e)
37    }
38}
39
40impl From<std::io::Error> for ZipError {
41    fn from(e: std::io::Error) -> Self {
42        ZipError::Io(e)
43    }
44}
45
46pub struct ZipEntry {
47    pub filename: String,
48    pub uncompressed_size: u32,
49    pub data: Vec<u8>,
50}
51
52pub struct MuyZipido {
53    response: Option<reqwest::blocking::Response>,
54    chunk_size: usize,
55    buffer: Vec<u8>,
56    offset: usize,
57    finished: bool,
58    progress_bar: Option<ProgressBar>,
59}
60
61impl MuyZipido {
62    pub fn new(url: &str, chunk_size: usize) -> Result<Self, ZipError> {
63        let response = reqwest::blocking::get(url)?;
64
65        if !response.status().is_success() {
66            return Err(ZipError::Http(response.error_for_status().unwrap_err()));
67        }
68
69        Ok(Self {
70            response: Some(response),
71            chunk_size,
72            buffer: Vec::new(),
73            offset: 0,
74            finished: false,
75            progress_bar: None,
76        })
77    }
78
79    pub fn with_progress(
80        mut self,
81        style: progress_bar::Style,
82        color: progress_bar::Colour,
83    ) -> Self {
84        let content_length = if let Some(response) = &self.response {
85            response
86                .headers()
87                .get("content-length")
88                .and_then(|value| value.to_str().ok())
89                .and_then(|s| s.parse::<usize>().ok())
90        } else {
91            None
92        };
93
94        let progress_bar = ProgressBar::new(content_length)
95            .with_description("Downloading ZIP".to_string())
96            .with_style(style)
97            .with_color(color);
98        self.progress_bar = Some(progress_bar);
99        self
100    }
101
102    fn read_exact(&mut self, size: usize) -> Result<Vec<u8>, ZipError> {
103        while self.buffer.len() < size {
104            if let Some(response) = &mut self.response {
105                let mut chunk = vec![0u8; self.chunk_size];
106                let bytes_read = response.read(&mut chunk)?;
107
108                if bytes_read == 0 {
109                    return Err(ZipError::UnexpectedEof);
110                }
111
112                chunk.truncate(bytes_read);
113                self.buffer.extend_from_slice(&chunk);
114
115                if let Some(ref mut progress_bar) = self.progress_bar {
116                    progress_bar.update(bytes_read);
117                }
118            } else {
119                return Err(ZipError::UnexpectedEof);
120            }
121        }
122
123        let data = self.buffer[..size].to_vec();
124        self.buffer.drain(..size);
125        self.offset += size;
126
127        Ok(data)
128    }
129
130    fn process_with_descriptor(&mut self, compression: u16) -> Result<Vec<u8>, ZipError> {
131        const DATA_DESC_SIG: [u8; 4] = [0x50, 0x4b, 0x07, 0x08];
132
133        let mut data = Vec::new();
134        let mut sig_buffer: CircularBuffer<u8> = CircularBuffer::new(4);
135
136        if compression == 8 {
137            let mut compressed_data = Vec::new();
138
139            loop {
140                let byte = self.read_exact(1)?[0];
141                compressed_data.push(byte);
142                sig_buffer.write(byte);
143
144                if sig_buffer.len() >= 4 {
145                    let last_4 = sig_buffer.get_last_n(4);
146                    if last_4.as_slice() == DATA_DESC_SIG {
147                        compressed_data.truncate(compressed_data.len() - 4);
148
149                        let mut decoder = DeflateDecoder::new(&compressed_data[..]);
150                        decoder.read_to_end(&mut data)?;
151
152                        let _crc = self.read_exact(4)?;
153                        let _compressed_size = self.read_exact(4)?;
154                        let _uncompressed_size = self.read_exact(4)?;
155
156                        break;
157                    }
158                }
159
160                if compressed_data.len() > 100_000_000 {
161                    return Err(ZipError::Decompression(
162                        "Data descriptor not found within reasonable limit".to_string(),
163                    ));
164                }
165            }
166        } else if compression == 0 {
167            loop {
168                let byte = self.read_exact(1)?[0];
169                data.push(byte);
170                sig_buffer.write(byte);
171
172                if sig_buffer.len() >= 4 {
173                    let last_4 = sig_buffer.get_last_n(4);
174                    if last_4.as_slice() == DATA_DESC_SIG {
175                        data.truncate(data.len() - 4);
176
177                        let _crc = self.read_exact(4)?;
178                        let _compressed_size = self.read_exact(4)?;
179                        let _uncompressed_size = self.read_exact(4)?;
180
181                        break;
182                    }
183                }
184
185                if data.len() > 100_000_000 {
186                    return Err(ZipError::Decompression(
187                        "Data descriptor not found within reasonable limit".to_string(),
188                    ));
189                }
190            }
191        } else {
192            return Err(ZipError::Decompression(format!(
193                "Unsupported compression method: {}",
194                compression
195            )));
196        }
197
198        Ok(data)
199    }
200
201    fn process_next_entry(&mut self) -> Result<Option<ZipEntry>, ZipError> {
202        const LOCAL_FILE_HEADER_SIG: &[u8] = b"PK\x03\x04";
203        const CENTRAL_DIR_SIG: &[u8] = b"PK\x01\x02";
204        const END_CENTRAL_DIR_SIG: &[u8] = b"PK\x05\x06";
205
206        if self.finished {
207            return Ok(None);
208        }
209
210        let sig = self.read_exact(4)?;
211
212        if sig == CENTRAL_DIR_SIG || sig == END_CENTRAL_DIR_SIG {
213            println!("Reached end of local file entries");
214            self.finished = true;
215            return Ok(None);
216        }
217
218        if sig != LOCAL_FILE_HEADER_SIG {
219            let mut hex_string = String::with_capacity(sig.len() * 2);
220            for b in &sig {
221                hex_string.push_str(&format!("{:02x}", b));
222            }
223            return Err(ZipError::InvalidSignature(hex_string));
224        }
225
226        let header_data = self.read_exact(26)?;
227        let _version = u16::from_le_bytes([header_data[0], header_data[1]]);
228        let flags = u16::from_le_bytes([header_data[2], header_data[3]]);
229        let compression = u16::from_le_bytes([header_data[4], header_data[5]]);
230        let _mod_time = u16::from_le_bytes([header_data[6], header_data[7]]);
231        let _mod_date = u16::from_le_bytes([header_data[8], header_data[9]]);
232        let _crc32 = u32::from_le_bytes([
233            header_data[10],
234            header_data[11],
235            header_data[12],
236            header_data[13],
237        ]);
238        let compressed_size = u32::from_le_bytes([
239            header_data[14],
240            header_data[15],
241            header_data[16],
242            header_data[17],
243        ]);
244        let uncompressed_size = u32::from_le_bytes([
245            header_data[18],
246            header_data[19],
247            header_data[20],
248            header_data[21],
249        ]);
250        let filename_len = u16::from_le_bytes([header_data[22], header_data[23]]);
251        let extra_len = u16::from_le_bytes([header_data[24], header_data[25]]);
252
253        let filename_bytes = self.read_exact(filename_len as usize)?;
254        let filename = String::from_utf8_lossy(&filename_bytes).to_string();
255        let _extra_field = self.read_exact(extra_len as usize)?;
256
257        let has_data_descriptor = (flags & 0x08) != 0;
258
259        println!("\nProcessing: {}", filename);
260        println!("  Compression: {} (0=none, 8=deflate)", compression);
261
262        let data = if !has_data_descriptor && compressed_size > 0 {
263            let compressed_data = self.read_exact(compressed_size as usize)?;
264
265            match compression {
266                0 => compressed_data,
267                8 => {
268                    let mut decoder = DeflateDecoder::new(&compressed_data[..]);
269                    let mut decompressed = Vec::new();
270                    decoder.read_to_end(&mut decompressed)?;
271                    decompressed
272                }
273                _ => {
274                    return Err(ZipError::Decompression(format!(
275                        "Unsupported compression method: {}",
276                        compression
277                    )));
278                }
279            }
280        } else if has_data_descriptor {
281            println!("  Streaming with data descriptor...");
282            self.process_with_descriptor(compression)?
283        } else {
284            Vec::new()
285        };
286
287        println!("  Processed {} bytes", data.len());
288
289        Ok(Some(ZipEntry {
290            filename,
291            uncompressed_size,
292            data,
293        }))
294    }
295}
296
297impl Drop for MuyZipido {
298    fn drop(&mut self) {
299        if let Some(ref mut progress_bar) = self.progress_bar {
300            progress_bar.finish();
301        }
302    }
303}
304
305impl Iterator for MuyZipido {
306    type Item = Result<ZipEntry, ZipError>;
307
308    fn next(&mut self) -> Option<Self::Item> {
309        match self.process_next_entry() {
310            Ok(Some(entry)) => Some(Ok(entry)),
311            Ok(None) => None,
312            Err(e) => {
313                self.finished = true;
314                Some(Err(e))
315            }
316        }
317    }
318}