pkcrack 0.1.0

//! ZIP file processing for pkcrack

use crate::types::*;
use crate::error::{PkCrackError, Result};
use crate::crypto::get_key_manager;
use std::io::{Read, Write, Seek, SeekFrom};
use std::fs::File;
use std::path::Path;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

/// ZIP file processor
pub struct ZipProcessor {
    key_manager: crate::crypto::KeyManager,
}

impl ZipProcessor {
    /// Create new ZIP processor
    pub fn new() -> Result<Self> {
        Ok(Self {
            key_manager: get_key_manager()?.clone(),
        })
    }

    /// Extract file from ZIP archive
    pub fn extract_file(&self, zip_path: &Path, filename: &str, case_sensitive: bool) -> Result<(Vec<u8>, LocalFileHeader)> {
        let mut file = File::open(zip_path)
            .map_err(|e| PkCrackError::FileNotFound(format!("{}: {}", zip_path.display(), e)))?;

        // Find the file in the ZIP
        let (local_header, file_data) = self.find_file_in_zip(&mut file, filename, case_sensitive)?;

        Ok((file_data, local_header))
    }

    /// Find file in ZIP and return its header and data
    fn find_file_in_zip(&self, file: &mut File, filename: &str, case_sensitive: bool) -> Result<(LocalFileHeader, Vec<u8>)> {
        // First try to find via local headers (simple ZIP)
        if let Ok((header, data)) = self.scan_local_headers(file, filename, case_sensitive) {
            return Ok((header, data));
        }

        // If not found, try to parse central directory
        file.seek(SeekFrom::Start(0))?;
        self.scan_central_directory(file, filename, case_sensitive)
    }

    /// Scan local file headers to find target file
    fn scan_local_headers(&self, file: &mut File, filename: &str, case_sensitive: bool) -> Result<(LocalFileHeader, Vec<u8>)> {
        file.seek(SeekFrom::Start(0))?;

        loop {
            let pos = file.stream_position()?;

            // Read signature
            let signature = file.read_u32::<LittleEndian>()?;
            if signature != LocalFileHeader::SIGNATURE {
                if pos == 0 {
                    return Err(PkCrackError::InvalidZip("Not a ZIP file".to_string()));
                } else {
                    break; // End of local headers
                }
            }

            // Read local header
            let header = self.read_local_file_header(file)?;

            // Read filename
            let mut filename_bytes = vec![0u8; header.file_name_length as usize];
            file.read_exact(&mut filename_bytes)?;
            let file_filename = String::from_utf8_lossy(&filename_bytes);

            // Skip extra field
            if header.extra_field_length > 0 {
                file.seek(SeekFrom::Current(header.extra_field_length as i64))?;
            }

            // Check if this is our target file
            let matches = if case_sensitive {
                file_filename == filename
            } else {
                file_filename.to_lowercase() == filename.to_lowercase()
            };

            if matches {
                // Read file data
                let mut data = vec![0u8; header.compressed_size as usize];
                file.read_exact(&mut data)?;

                return Ok((header, data));
            } else {
                // Skip file data
                file.seek(SeekFrom::Current(header.compressed_size as i64))?;

                // Skip data descriptor if present
                if header.general_purpose & 0x8 != 0 {
                    // Data descriptor present - skip it
                    // Try to read signature first
                    let pos = file.stream_position()?;
                    let possible_sig = file.read_u32::<LittleEndian>().unwrap_or(0);
                    if possible_sig == DataDescriptor::SIGNATURE {
                        file.seek(SeekFrom::Current(12))?; // Skip rest of descriptor
                    } else {
                        // No signature, seek back and skip fixed size
                        file.seek(SeekFrom::Start(pos))?;
                        file.seek(SeekFrom::Current(12))?;
                    }
                }
            }
        }

        Err(PkCrackError::FileNotFoundInZip(filename.to_string()))
    }

    /// Scan central directory to find file
    fn scan_central_directory(&self, file: &mut File, filename: &str, case_sensitive: bool) -> Result<(LocalFileHeader, Vec<u8>)> {
        // Find end of central directory record
        let (end_central_dir, _) = self.find_end_central_directory(file)?;

        // Seek to central directory
        file.seek(SeekFrom::Start(end_central_dir.central_dir_offset as u64))?;

        // Scan central directory entries
        for _ in 0..end_central_dir.total_entries {
            let signature = file.read_u32::<LittleEndian>()?;
            if signature != CentralDirHeader::SIGNATURE {
                return Err(PkCrackError::InvalidZip("Invalid central directory signature".to_string()));
            }

            let central_header = self.read_central_dir_header(file)?;

            // Read filename
            let mut filename_bytes = vec![0u8; central_header.file_name_length as usize];
            file.read_exact(&mut filename_bytes)?;
            let file_filename = String::from_utf8_lossy(&filename_bytes);

            // Skip extra field and comment
            file.seek(SeekFrom::Current(
                (central_header.extra_field_length + central_header.file_comment_length) as i64
            ))?;

            // Check if this is our target file
            let matches = if case_sensitive {
                file_filename == filename
            } else {
                file_filename.to_lowercase() == filename.to_lowercase()
            };

            if matches {
                // Seek to local header
                file.seek(SeekFrom::Start(central_header.local_header_offset as u64))?;

                // Read local header and file data
                let local_header = self.read_local_file_header(file)?;

                // Skip filename and extra field
                file.seek(SeekFrom::Current(
                    (local_header.file_name_length + local_header.extra_field_length) as i64
                ))?;

                // Read file data
                let mut data = vec![0u8; local_header.compressed_size as usize];
                file.read_exact(&mut data)?;

                return Ok((local_header, data));
            }
        }

        Err(PkCrackError::FileNotFoundInZip(filename.to_string()))
    }

    /// Find end of central directory record
    fn find_end_central_directory(&self, file: &mut File) -> Result<(EndCentralDirRecord, u64)> {
        // Start from end of file and search backwards
        let file_size = file.seek(SeekFrom::End(0))?;
        let max_comment_length = 65535 + 22; // Max comment + header size
        let search_start = if file_size > max_comment_length {
            file_size - max_comment_length
        } else {
            0
        };

        file.seek(SeekFrom::Start(search_start))?;

        let mut buffer = vec![0u8; (file_size - search_start) as usize];
        file.read_exact(&mut buffer)?;

        // Search for signature
        for i in (0..buffer.len() - 4).rev() {
            if buffer[i..i+4] == [0x50, 0x4b, 0x05, 0x06] { // 0x06054b50
                file.seek(SeekFrom::Start(search_start + i as u64))?;
                let end_record = self.read_end_central_dir(file)?;
                return Ok((end_record, search_start + i as u64));
            }
        }

        Err(PkCrackError::InvalidZip("End of central directory not found".to_string()))
    }

    /// Read local file header
    fn read_local_file_header(&self, file: &mut File) -> Result<LocalFileHeader> {
        Ok(LocalFileHeader {
            signature: LocalFileHeader::SIGNATURE,
            version_needed: file.read_u16::<LittleEndian>()?,
            general_purpose: file.read_u16::<LittleEndian>()?,
            compression_method: file.read_u16::<LittleEndian>()?,
            last_mod_time: file.read_u16::<LittleEndian>()?,
            last_mod_date: file.read_u16::<LittleEndian>()?,
            crc32: file.read_u32::<LittleEndian>()?,
            compressed_size: file.read_u32::<LittleEndian>()?,
            uncompressed_size: file.read_u32::<LittleEndian>()?,
            file_name_length: file.read_u16::<LittleEndian>()?,
            extra_field_length: file.read_u16::<LittleEndian>()?,
        })
    }

    /// Read central directory header
    fn read_central_dir_header(&self, file: &mut File) -> Result<CentralDirHeader> {
        Ok(CentralDirHeader {
            signature: CentralDirHeader::SIGNATURE,
            version_made_by: file.read_u16::<LittleEndian>()?,
            version_needed: file.read_u16::<LittleEndian>()?,
            general_purpose: file.read_u16::<LittleEndian>()?,
            compression_method: file.read_u16::<LittleEndian>()?,
            last_mod_time: file.read_u16::<LittleEndian>()?,
            last_mod_date: file.read_u16::<LittleEndian>()?,
            crc32: file.read_u32::<LittleEndian>()?,
            compressed_size: file.read_u32::<LittleEndian>()?,
            uncompressed_size: file.read_u32::<LittleEndian>()?,
            file_name_length: file.read_u16::<LittleEndian>()?,
            extra_field_length: file.read_u16::<LittleEndian>()?,
            file_comment_length: file.read_u16::<LittleEndian>()?,
            disk_number_start: file.read_u16::<LittleEndian>()?,
            internal_attributes: file.read_u16::<LittleEndian>()?,
            external_attributes: file.read_u32::<LittleEndian>()?,
            local_header_offset: file.read_u32::<LittleEndian>()?,
        })
    }

    /// Read end of central directory record
    fn read_end_central_dir(&self, file: &mut File) -> Result<EndCentralDirRecord> {
        Ok(EndCentralDirRecord {
            signature: EndCentralDirRecord::SIGNATURE,
            disk_number: file.read_u16::<LittleEndian>()?,
            central_dir_disk: file.read_u16::<LittleEndian>()?,
            entries_on_disk: file.read_u16::<LittleEndian>()?,
            total_entries: file.read_u16::<LittleEndian>()?,
            central_dir_size: file.read_u32::<LittleEndian>()?,
            central_dir_offset: file.read_u32::<LittleEndian>()?,
            comment_length: file.read_u16::<LittleEndian>()?,
        })
    }

    /// Decrypt ZIP file to output file
    pub fn decrypt_zip_to_file(&self, zip_path: &Path, filename: &str, output_path: &str, key_state: &mut KeyState) -> Result<()> {
        // Extract file from ZIP
        let (mut encrypted_data, header) = self.extract_file(zip_path, filename, false)?;

        // Decrypt the data
        self.key_manager.decrypt(key_state, &mut encrypted_data);

        // Create output ZIP file
        self.create_decrypted_zip(output_path, &header, filename, &encrypted_data)?;

        Ok(())
    }

    /// Create decrypted ZIP file
    fn create_decrypted_zip(&self, output_path: &str, header: &LocalFileHeader, filename: &str, decrypted_data: &[u8]) -> Result<()> {
        let mut output_file = File::create(output_path)?;

        // Write local file header
        self.write_local_file_header(&mut output_file, header, filename, decrypted_data)?;

        // Write file data
        output_file.write_all(decrypted_data)?;

        // Write data descriptor if needed
        if header.general_purpose & 0x8 != 0 {
            self.write_data_descriptor(&mut output_file, decrypted_data)?;
        }

        // Write end of central directory record
        self.write_end_central_dir(&mut output_file, filename, decrypted_data)?;

        Ok(())
    }

    /// Write local file header to output
    fn write_local_file_header(&self, file: &mut File, header: &LocalFileHeader, filename: &str, data: &[u8]) -> Result<()> {
        file.write_u32::<LittleEndian>(LocalFileHeader::SIGNATURE)?;
        file.write_u16::<LittleEndian>(header.version_needed)?;
        file.write_u16::<LittleEndian>(header.general_purpose & !0x1)?; // Clear encryption flag
        file.write_u16::<LittleEndian>(header.compression_method)?;
        file.write_u16::<LittleEndian>(header.last_mod_time)?;
        file.write_u16::<LittleEndian>(header.last_mod_date)?;
        file.write_u32::<LittleEndian>(crate::crc::crc32_buffer_fast(0xFFFFFFFF, data) ^ 0xFFFFFFFF)?; // Calculate real CRC
        file.write_u32::<LittleEndian>(data.len() as u32)?; // Uncompressed size (same as original)
        file.write_u32::<LittleEndian>(data.len() as u32)?; // Compressed size (same as original)
        file.write_u16::<LittleEndian>(filename.len() as u16)?;
        file.write_u16::<LittleEndian>(0)?; // No extra field

        // Write filename
        file.write_all(filename.as_bytes())?;

        Ok(())
    }

    /// Write data descriptor
    fn write_data_descriptor(&self, file: &mut File, data: &[u8]) -> Result<()> {
        file.write_u32::<LittleEndian>(DataDescriptor::SIGNATURE)?;
        file.write_u32::<LittleEndian>(crate::crc::crc32_buffer_fast(0xFFFFFFFF, data) ^ 0xFFFFFFFF)?;
        file.write_u32::<LittleEndian>(data.len() as u32)?;
        file.write_u32::<LittleEndian>(data.len() as u32)?;

        Ok(())
    }

    /// Write end of central directory record
    fn write_end_central_dir(&self, file: &mut File, filename: &str, data: &[u8]) -> Result<()> {
        let central_dir_size = 46 + filename.len() as u32; // Central dir header + filename
        let central_dir_offset = 30 + filename.len() as u32 + data.len() as u32; // Local header + filename + data

        file.write_u32::<LittleEndian>(EndCentralDirRecord::SIGNATURE)?;
        file.write_u16::<LittleEndian>(0)?; // Disk number
        file.write_u16::<LittleEndian>(0)?; // Central dir disk
        file.write_u16::<LittleEndian>(1)?; // Entries on disk
        file.write_u16::<LittleEndian>(1)?; // Total entries
        file.write_u32::<LittleEndian>(central_dir_size)?;
        file.write_u32::<LittleEndian>(central_dir_offset)?;
        file.write_u16::<LittleEndian>(0)?; // Comment length

        Ok(())
    }

    /// Check if file in ZIP is encrypted
    pub fn is_file_encrypted(&self, zip_path: &Path, filename: &str) -> Result<bool> {
        let mut file = File::open(zip_path)?;
        let (header, _) = self.find_file_in_zip(&mut file, filename, false)?;
        Ok(header.general_purpose & 0x1 != 0)
    }

    /// List files in ZIP archive
    pub fn list_files(&self, zip_path: &Path) -> Result<Vec<String>> {
        let mut files = Vec::new();
        let mut file = File::open(zip_path)?;

        // Find end of central directory
        let (end_central_dir, _) = self.find_end_central_directory(&mut file)?;

        // Seek to central directory
        file.seek(SeekFrom::Start(end_central_dir.central_dir_offset as u64))?;

        // Scan central directory entries
        for _ in 0..end_central_dir.total_entries {
            let signature = file.read_u32::<LittleEndian>()?;
            if signature != CentralDirHeader::SIGNATURE {
                break;
            }

            let central_header = self.read_central_dir_header(&mut file)?;

            // Read filename
            let mut filename_bytes = vec![0u8; central_header.file_name_length as usize];
            file.read_exact(&mut filename_bytes)?;
            let filename = String::from_utf8_lossy(&filename_bytes).to_string();

            // Skip extra field and comment
            file.seek(SeekFrom::Current(
                (central_header.extra_field_length + central_header.file_comment_length) as i64
            ))?;

            files.push(filename);
        }

        Ok(files)
    }
}

impl Default for ZipProcessor {
    fn default() -> Self {
        Self::new().expect("Failed to create ZipProcessor")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    #[test]
    fn test_zip_processor_creation() {
        let processor = ZipProcessor::new();
        assert!(processor.is_ok());
    }

    #[test]
    fn test_write_local_file_header() {
        let processor = ZipProcessor::new().unwrap();
        let header = LocalFileHeader {
            signature: LocalFileHeader::SIGNATURE,
            version_needed: 20,
            general_purpose: 0,
            compression_method: 0,
            last_mod_time: 0,
            last_mod_date: 0,
            crc32: 0,
            compressed_size: 12,
            uncompressed_size: 12,
            file_name_length: 8,
            extra_field_length: 0,
        };

        let filename = "test.txt";
        let data = b"Hello World";

 
    }

    #[test]
    fn test_is_file_encrypted() {
        let processor = ZipProcessor::new().unwrap();

        // Test with non-existent file
        let path = PathBuf::from("nonexistent.zip");
        let result = processor.is_file_encrypted(&path, "test.txt");
        assert!(result.is_err());
    }
}