pmv_encryption_rs 1.0.0

// Read stream

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

use byteorder::{BigEndian, ByteOrder};

use crate::{BlockFileOpenError, BlockFileReadError, decrypt};

/// A stream to read an encrypted large file
/// using blocks for fast skip
pub struct FileBlockEncryptReadStream {
    /// File descriptor
    file: File,

    /// File size in bytes
    file_size: u64,

    /// Block size in bytes
    block_size: u64,

    /// Block count
    block_count: u64,

    /// Encryption key
    key: Vec<u8>,

    /// Cursor position
    cur_pos: u64,

    /// Current block the cursor is reading
    cur_block: Option<u64>,

    /// Buffer to store the current block (Encrypted)
    cur_block_encrypted: Vec<u8>,

    /// Buffer to store the current block (Decrypted)
    cur_block_decrypted: Vec<u8>,
}

impl FileBlockEncryptReadStream {
    /// Creates a new instance of FileBlockEncryptReadStream
    /// Opens the file and reads the header
    ///
    /// Parameters:
    ///  - `file_path` - Path to the file to open
    ///  - `key` - The encryption key
    ///
    /// Returns an instance of FileBlockEncryptReadStream, or an error
    pub fn new<P>(
        file_path: P,
        key: Vec<u8>,
    ) -> Result<FileBlockEncryptReadStream, BlockFileOpenError>
    where
        P: AsRef<Path>,
    {
        let mut file = File::open(file_path)?;

        // Read file size

        let mut buf: [u8; 8] = [0; 8];

        file.read_exact(&mut buf)?;

        let file_size = BigEndian::read_u64(&buf);

        // Read block size

        file.read_exact(&mut buf)?;

        let block_size = BigEndian::read_u64(&buf);

        if block_size == 0 {
            return Err(BlockFileOpenError::InvalidBlockSize);
        }

        let mut block_count = file_size / block_size;

        if file_size % block_size != 0 {
            block_count += 1;
        }

        Ok(FileBlockEncryptReadStream {
            file,
            key,
            file_size,
            block_size,
            block_count,
            cur_block: None,
            cur_pos: 0,
            cur_block_encrypted: Vec::new(),
            cur_block_decrypted: Vec::new(),
        })
    }

    /// Gets the size of the file
    pub fn get_file_size(&self) -> u64 {
        self.file_size
    }

    /// Gets the size of the blocks in the file
    pub fn get_block_size(&self) -> u64 {
        self.block_size
    }

    /// Gets the total number of blocks in the file
    pub fn get_block_count(&self) -> u64 {
        self.block_count
    }

    /// Gets the current cursor position
    pub fn get_cursor(&self) -> u64 {
        self.cur_pos
    }

    /// Fetches a block
    fn fetch_block(&mut self, block_num: u64) -> Result<(), BlockFileReadError> {
        if block_num >= self.block_count {
            return Err(BlockFileReadError::OutOfBounds);
        }

        // Read block index entry

        let mut buf_block_index: [u8; 16] = [0; 16];

        self.file.seek(io::SeekFrom::Start(16 + block_num * 16))?;
        self.file.read_exact(&mut buf_block_index)?;

        let block_start = BigEndian::read_u64(&buf_block_index[0..8]);
        let block_length = BigEndian::read_u64(&buf_block_index[8..16]);

        // Seek to the data and read it

        self.file.seek(io::SeekFrom::Start(block_start))?;

        self.cur_block_encrypted.resize(block_length as usize, 0);

        self.file.read_exact(&mut self.cur_block_encrypted)?;

        // Decrypt

        self.cur_block_decrypted = decrypt(&self.cur_block_encrypted, &self.key)?;

        // Update pointer

        self.cur_block = Some(block_num);

        Ok(())
    }

    /// Seeks the file, changing the position of the cursor
    ///
    /// Parameters:
    ///
    ///  - `pos` - The position in the file to seek to
    ///
    /// Returns the new position of the cursor, or an error
    pub fn seek(&mut self, pos: SeekFrom) -> Result<u64, BlockFileReadError> {
        let new_pos = match pos {
            SeekFrom::Start(o) => o,
            SeekFrom::End(o) => (self.file_size as i64).wrapping_add(o) as u64,
            SeekFrom::Current(o) => (self.cur_pos as i64).wrapping_add(o) as u64,
        };

        if new_pos > self.file_size {
            return Err(BlockFileReadError::OutOfBounds);
        }

        self.cur_pos = new_pos;

        Ok(new_pos)
    }

    /// Reads data into a buffer
    ///
    /// Parameters:
    ///
    ///  - `buf` - The buffer for the data to be placed into
    ///
    /// Returns the number of bytes read, or an error
    /// When the files does not contain more bytes, the error will be BlockFileReadError::EndOfFile
    pub fn read(&mut self, buf: &mut [u8]) -> Result<usize, BlockFileReadError> {
        if self.cur_pos >= self.file_size {
            return Err(BlockFileReadError::EndOfFile);
        }

        let mut filled_length: usize = 0;

        while filled_length < buf.len() && self.cur_pos < self.file_size {
            // Compute block index and fetch it

            let block_index = self.cur_pos / self.block_size;
            let block_offset = (self.cur_pos % self.block_size) as usize;

            if self.cur_block.is_none() || self.cur_block.unwrap() != block_index {
                self.fetch_block(block_index)?;
            }

            // Compute total bytes to copy from this block

            let block_length = self.cur_block_decrypted.len();

            if block_offset > block_length {
                return Err(BlockFileReadError::OutOfBounds);
            }

            let mut bytes_to_copy = block_length - block_offset;

            let bytes_can_fit = buf.len() - filled_length;

            if bytes_to_copy > bytes_can_fit {
                bytes_to_copy = bytes_can_fit;
            }

            // Copy bytes

            buf[filled_length..filled_length + bytes_to_copy].copy_from_slice(
                &self.cur_block_decrypted[block_offset..block_offset + bytes_to_copy],
            );

            filled_length += bytes_to_copy;

            // Update cursor

            self.cur_pos += bytes_to_copy as u64;
        }

        Ok(filled_length)
    }

    /// Closes the file and drops the instance of the stream
    pub fn close(self) {}
}