psarc2 0.1.0

//! Archive reading functionality.

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

use crate::{
    PlaystationArchive,
    compression::CompressionType,
    error::{Error, Result},
    file::{BasicFileEntry as _, FileEntry},
    path::ArchiveFlags,
    toc::{DecryptionKey, TableOfContent},
};

/// Read configuration options.
#[derive(Default)]
pub struct Config {
    /// Keys for decrypting the header.
    pub decryption_key: Option<DecryptionKey>,
}

/// How big each data block is in bytes.
#[derive(Debug, Clone, Copy)]
pub(crate) enum BlockSize {
    /// 2 bytes.
    U16,
    /// 3 bytes.
    U24,
    /// 4 bytes.
    U32,
}

impl BlockSize {
    /// Convert the blocksize to it's number representation.
    pub(crate) const fn as_u32(self) -> u32 {
        match self {
            Self::U16 => 0x0001_0000,
            Self::U24 => 0x0100_0000,
            Self::U32 => u32::MAX,
        }
    }

    /// Read a block entry based on the type used.
    pub(crate) fn read<R>(self, reader: &mut R) -> Result<u32>
    where
        R: Read + Seek,
    {
        match self {
            // Read as u16 then cast
            Self::U16 => Ok(u32::from(read_u16(reader)?)),
            // Read as u64 fitted, then try to cast (should always succeed)
            Self::U24 => u32::try_from(read_unsigned::<3>(reader)?)
                .map_err(|_| Error::Corrupt("3 bytes got parsed as more")),
            // Read directly as u32
            Self::U32 => read_u32(reader),
        }
    }
}

impl TryFrom<u32> for BlockSize {
    type Error = Error;

    fn try_from(value: u32) -> Result<Self> {
        match value {
            0x0001_0000 => Ok(Self::U16),
            0x0100_0000 => Ok(Self::U24),
            u32::MAX => Ok(Self::U32),
            _ => Err(Error::Corrupt("unregular block size")),
        }
    }
}

impl<R: Read + Seek> PlaystationArchive<R> {
    /// Read a Playstation archive file using the default options, collecting the files it contains.
    ///
    /// This uses the central directory record of the `.psarc` file, known as the manifest file, and ignores local file headers.
    ///
    /// # Errors
    ///
    /// - When archive file cannot be read.
    /// - When file is corrupt.
    ///
    /// # Example
    ///
    /// ```
    /// // Read a file from disk
    /// let file = File::open("file.psarc")?;
    ///
    /// // Open it as an archive
    /// let archive = PlaystationArchive::new(file)?;
    /// ```
    pub fn new(reader: R) -> Result<Self> {
        Self::with_config(Config::default(), reader)
    }

    /// Read a Playstation archive file, collecting the files it contains.
    ///
    /// This uses the central directory record of the `.psarc` file, known as the manifest file, and ignores local file headers.
    ///
    /// # Errors
    ///
    /// - When archive file cannot be read.
    /// - When file is corrupt.
    pub fn with_config(config: Config, mut reader: R) -> Result<Self> {
        // Check the magic bytes
        let mut magic = [0_u8; 4];
        reader.read_exact(&mut magic)?;
        if magic != *b"PSAR" {
            return Err(Error::UnrecognizedFile);
        }

        // Check the version
        let major_version = read_u16(&mut reader)?;
        let minor_version = read_u16(&mut reader)?;
        if major_version != 1 || minor_version != 4 {
            return Err(Error::UnsupportedVersion);
        }

        // Read the compression type
        let compression_type: CompressionType = read_u32(&mut reader)?.try_into()?;

        // Read the table of contents
        let table_of_content = TableOfContent::read(&mut reader)?;

        // Read the block size
        let block_size: BlockSize = read_u32(&mut reader)?.try_into()?;

        // Read the archive flags
        let archive_flags: ArchiveFlags = read_u32(&mut reader)?.try_into()?;

        // Get all entries from the table of content
        let table_of_content_entries =
            table_of_content.read_entries(&mut reader, archive_flags, config.decryption_key)?;

        // Read the blocks
        let blocks_amount = table_of_content.blocks_amount();

        // Parse the block sizes
        let block_sizes = (0..blocks_amount)
            .map(|_| block_size.read(&mut reader))
            .collect::<Result<Vec<_>>>()?;

        // Read the manifest
        let manifest_entry = table_of_content_entries
            .first()
            .ok_or(Error::FileAtIndexDoesNotExist(0))?;
        let mut manifest_bytes = vec![0; manifest_entry.input_size()];
        let mut manifest_writer = Cursor::new(&mut manifest_bytes);
        compression_type.decompress(
            &mut reader,
            &mut manifest_writer,
            manifest_entry,
            block_size,
            &block_sizes,
        )?;

        // Convert to string
        let manifest_string = String::from_utf8(manifest_bytes)
            .map_err(|_| Error::Corrupt("Manifest is not valid UTF-8"))?;

        // Create file entries
        let file_entries = std::iter::once("manifest.txt")
            .chain(manifest_string.lines())
            .zip(table_of_content_entries)
            .map(|(path, table_of_content_entry)| {
                FileEntry::from_table_of_content(table_of_content_entry, path)
            })
            .collect();

        Ok(Self {
            compression_type,
            file_entries,
            block_size,
            reader,
            block_sizes,
        })
    }

    /// Extract the archive into a directory, throwing an error if they already exist.
    ///
    /// # Errors
    ///
    /// - When any of the files already exists.
    /// - When extracting files fails.
    /// - When directory can't be created.
    /// - When directory doesn't have the correct permissions.
    /// - When files can't be written.
    pub fn extract<P>(&mut self, directory: P) -> Result<()>
    where
        P: AsRef<Path>,
    {
        let directory = directory.as_ref();

        // Create the directory if it doesn't exist
        std::fs::create_dir_all(directory)?;

        // Canonicalize the directory to append the rest
        let directory = directory.canonicalize()?;

        // Copy the entries so we can use a mutable reference
        let file_entries = self
            .file_entries
            .iter()
            .skip(1)
            .cloned()
            .collect::<Vec<_>>();

        // Read each file, skipping the manifest
        for file_entry in file_entries {
            // Extract each file
            let file_path = directory.join(&file_entry.path);

            // Create the directory
            if let Some(file_dir) = file_path.parent() {
                std::fs::create_dir_all(file_dir)?;
            }

            // Extract the file
            let mut file = File::create_new(file_path)?;
            self.write_bytes_by_entry(&mut file, &file_entry)?;
        }

        Ok(())
    }
}

/// Read a single `u64` from any amount of bytes from the buffer.
#[inline]
pub(crate) fn read_unsigned<const N: usize>(reader: &mut impl Read) -> Result<usize> {
    // Fill a buffer for the whole number
    let mut buf = [0_u8; size_of::<usize>()];
    // Only read the required bytes
    reader.read_exact(&mut buf[const { size_of::<usize>() - N }..])?;

    Ok(usize::from_be_bytes(buf))
}

/// Read a single `u32` from the buffer.
#[inline]
pub(crate) fn read_u32<R>(reader: &mut R) -> Result<u32>
where
    R: Read,
{
    // Read the required bytes
    let mut buf = [0_u8; size_of::<u32>()];
    reader.read_exact(&mut buf)?;

    Ok(u32::from_be_bytes(buf))
}

/// Read a single `u16` from the buffer.
#[inline]
pub(crate) fn read_u16<R>(reader: &mut R) -> Result<u16>
where
    R: Read,
{
    // Read the required bytes
    let mut buf = [0_u8; size_of::<u16>()];
    reader.read_exact(&mut buf)?;

    Ok(u16::from_be_bytes(buf))
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use super::*;

    #[test]
    fn u64_bytes() {
        // Source bytes
        let test_value = 123_456_789_usize;
        let be_bytes = Cursor::new(test_value.to_be_bytes());

        // Test with different amounts
        let mut test = be_bytes.clone();
        assert_eq!(
            read_unsigned::<8>(&mut test).expect("Error reading"),
            test_value
        );

        let mut test = be_bytes.clone();
        test.seek_relative(1).expect("Unexpected EOB");
        assert_eq!(
            read_unsigned::<7>(&mut test).expect("Error reading"),
            test_value
        );

        let mut test = be_bytes.clone();
        test.seek_relative(2).expect("Unexpected EOB");
        assert_eq!(
            read_unsigned::<6>(&mut test).expect("Error reading"),
            test_value
        );

        let mut test = be_bytes;
        test.seek_relative(3).expect("Unexpected EOB");
        assert_eq!(
            read_unsigned::<5>(&mut test).expect("Error reading"),
            test_value
        );
    }
}