unarchiver.rs

use std::{
    fs::File,
    io::{BufWriter, Cursor, Read, Seek, SeekFrom, Write},
    path::Path,
};

use filetime::FileTime;

use crate::{
    codec::ArCodec,
    datamodel::{AllZeroExt, ArEntry, BlockPool, CompressionType, EntryType},
    read_write_extension::ReadExtTypes,
};

pub struct UnArchiver<T: Read + Seek> {
    codec: ArCodec,
    blocks_in: BlockPool<T>,
    metadata_compare: bool,
}

impl<T: Read + Seek> UnArchiver<T> {
    pub fn new(mut codec: ArCodec, blocks_in: BlockPool<T>) -> Self {
        // IndexTable is not always loaded for unarchiving, so use the BlockPool UID if it is not
        // set
        if codec.index_table.uid.is_all_zero() {
            codec.index_table.uid = blocks_in.uid;
        }

        Self {
            codec,
            blocks_in,
            metadata_compare: false,
        }
    }

    pub fn set_metadata_compare(&mut self, metadata_compare: bool) {
        self.metadata_compare = metadata_compare;
    }

    /// Check that all loaded files have the same UID
    ///
    /// Returns true if the UIDs are the same, false otherwise
    pub fn check_uids(&self) -> bool {
        self.codec.file_table.uid == self.codec.index_table.uid
            && self.codec.file_table.uid == self.blocks_in.uid
    }

    pub fn codec(&self) -> &ArCodec {
        &self.codec
    }

    pub fn entries_iter(&self) -> impl Iterator<Item = &ArEntry> {
        self.codec.file_table.entries.iter()
    }

    pub fn unarchive_to(&mut self, out_dir: impl AsRef<Path>) -> Result<(), std::io::Error> {
        // Swap the entries out of `self` for the loop to prevent borrowing issues
        let mut file_table_entries = Vec::new();
        std::mem::swap(&mut file_table_entries, &mut self.codec.file_table.entries);

        for entry in file_table_entries.iter() {
            // Removing the ':' fixes absolute paths in the archive on windows
            let entry_path = &entry.path.replace(':', "");
            let out_path = out_dir.as_ref().join(entry_path);

            if let Err(e) = self.unarchive_entry_to_file_path(entry, out_path) {
                println!("Error unarchiving entry: {} ({e})", entry.path);
            }
        }

        std::mem::swap(&mut file_table_entries, &mut self.codec.file_table.entries);

        Ok(())
    }

    /// Note: This will automatically create the parent directories and overwrite the file if it
    /// exists already
    pub fn unarchive_entry_to_file_path(
        &mut self,
        entry: &ArEntry,
        out_path: impl AsRef<Path>,
    ) -> Result<(), std::io::Error> {
        // If metadata compare is enabled, skip files that already exist and have the same mtime
        // and size
        if self.metadata_compare && out_path.as_ref().exists() {
            let orig_md = out_path.as_ref().metadata()?;
            let orig_mtime = filetime::FileTime::from_last_modification_time(&orig_md);
            let entry_mtime = filetime::FileTime::from_unix_time(
                entry.metadata.modified_unix_seconds,
                entry.metadata.modified_nanos,
            );

            if orig_md.len() == entry.metadata.file_size && orig_mtime == entry_mtime {
                return Ok(());
            }
        }

        if let Some(parent) = out_path.as_ref().parent() {
            std::fs::create_dir_all(parent)?;
        }

        if let Some(symlink_target) = entry.symlink_target.as_ref() {
            #[cfg(target_os = "windows")]
            {
                if out_path.as_ref().exists() {
                    std::fs::remove_file(out_path.as_ref())?;
                }

                let res = match entry.metadata.entry_and_os_type.entry_type() {
                    EntryType::SymbolicLinkFile => {
                        std::os::windows::fs::symlink_file(symlink_target, out_path.as_ref())
                    }
                    EntryType::SymbolicLinkDir => {
                        std::os::windows::fs::symlink_dir(symlink_target, out_path.as_ref())
                    }
                    _ => panic!("Trying to create symlink for non symlink entry"),
                };

                if let Err(e) = res {
                    // Windows error code for missing symlink privileges
                    match e.raw_os_error().unwrap_or(0) {
                        1314 => {
                            println!(
                                "Skipping symlink creation due to missing privileges: {}",
                                out_path.as_ref().display()
                            );
                            return Ok(());
                        }
                        _ => return Err(e),
                    }
                }
            }

            #[cfg(target_os = "linux")]
            {
                std::os::unix::fs::symlink(symlink_target, out_path.as_ref())?;
            }

            return Ok(());
        }

        if entry.metadata.entry_and_os_type.entry_type() == EntryType::Directory {
            if !out_path.as_ref().exists() {
                std::fs::create_dir(&out_path)?;
            }
            return Ok(());
        }

        let out_file = File::create(out_path.as_ref())?;
        let out_file = BufWriter::new(out_file);
        self.unarchive_entry_to_stream(entry, out_file)?;

        let last_modified = FileTime::from_unix_time(
            entry.metadata.modified_unix_seconds,
            entry.metadata.modified_nanos,
        );
        filetime::set_file_mtime(out_path.as_ref(), last_modified)?;

        Ok(())
    }

    pub fn unarchive_entry_to_stream(
        &mut self,
        entry: &ArEntry,
        mut stream_out: impl Write,
    ) -> Result<(), std::io::Error> {
        for start in entry.blocks.iter().copied() {
            self.blocks_in.seek(SeekFrom::Start(start))?;

            let block_len = self.blocks_in.read_exact_u32()?;
            self.codec.buffer.resize(block_len as usize, 0);

            let compression_type = self.blocks_in.read_exact_u8()?;
            let compression_type: CompressionType = compression_type.try_into()?;

            self.blocks_in.read_exact(&mut self.codec.buffer)?;

            match compression_type {
                CompressionType::Zstd => {
                    let mut decoder = zstd::Decoder::new(Cursor::new(&self.codec.buffer))?;
                    std::io::copy(&mut decoder, &mut stream_out)?;
                }
                CompressionType::Uncompressed => {
                    let mut cursor = Cursor::new(&self.codec.buffer);
                    std::io::copy(&mut cursor, &mut stream_out)?;
                }
            }
        }

        Ok(())
    }
}