pmv_encryption_rs 0.1.0

// Write stream

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

use byteorder::{BigEndian, ByteOrder};

use crate::{BlockWriteError, EncryptionMethod, encrypt};

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

    /// File size in bytes
    file_size: u64,

    // Written size
    written_size: u64,

    /// Block size in bytes
    block_size: u64,

    /// Block count
    block_count: u64,

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

    /// Current block being written
    current_write_index: u64,

    /// Position of the file to write the next block
    current_write_pt: u64,

    /// Write buffer
    write_buf: Vec<u8>,
}

impl FileBlockEncryptWriteStream {
    /// Creates a new instance of FileBlockEncryptWriteStream
    /// Creates the file and writes the header
    ///
    /// Parameters:
    ///  - `file_path` - Path to the file to create
    ///  - `key` - The encryption key
    ///  - `file_size` - The size of the file (in bytes)
    ///  - `block_size` - The max size for a block (in bytes)
    ///
    /// Return the instance of FileBlockEncryptWriteStream, or an error.
    pub fn new<P>(
        file_path: P,
        key: Vec<u8>,
        file_size: u64,
        block_size: u64,
    ) -> Result<FileBlockEncryptWriteStream, io::Error>
    where
        P: AsRef<Path>,
    {
        let mut block_count = file_size / block_size;

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

        let mut file = File::create(file_path)?;

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

        // Write file size

        BigEndian::write_u64(&mut buf, file_size);

        file.write_all(&buf)?;

        // Write block size

        BigEndian::write_u64(&mut buf, block_size);

        file.write_all(&buf)?;

        // Write default block index values

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

        for _i in 0..block_count {
            file.write_all(&buf_block_index)?;
        }

        // Return the struct

        Ok(FileBlockEncryptWriteStream {
            file,
            key,
            file_size,
            written_size: 0,
            block_size,
            block_count,
            current_write_index: 0,
            current_write_pt: 16 + (16 * block_count),
            write_buf: Vec::new(),
        })
    }

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

    /// Writes the current block to the file
    fn write_current_block(&mut self, encrypted_block: &[u8]) -> Result<(), BlockWriteError> {
        // Write block index entry

        self.file
            .seek(io::SeekFrom::Start(16 + self.current_write_index * 16))?;

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

        BigEndian::write_u64(&mut buf_block_index[0..8], self.current_write_pt); // Start pointer
        BigEndian::write_u64(&mut buf_block_index[8..16], encrypted_block.len() as u64); // Block length

        self.file.write_all(&buf_block_index)?;

        // Write block

        self.file.seek(io::SeekFrom::Start(self.current_write_pt))?;
        self.file.write_all(encrypted_block)?;

        // Update pointer

        self.current_write_index += 1;
        self.current_write_pt += encrypted_block.len() as u64;

        Ok(())
    }

    /// Writes data into the file
    /// The data will be automatically split into blocks
    /// and written to the file
    ///
    /// Parameters:
    ///
    ///  - `data` - Data to write
    ///
    /// Returns en empty tuple, or an error
    pub fn write(&mut self, data: &[u8]) -> Result<(), BlockWriteError> {
        if self.current_write_index >= self.block_count {
            return Err(BlockWriteError::ExceededFileSize);
        }

        self.written_size += data.len() as u64;

        self.write_buf.append(&mut data.to_vec());

        let block_size = self.block_size as usize;

        while self.write_buf.len() >= block_size {
            if self.current_write_index >= self.block_count {
                return Err(BlockWriteError::ExceededFileSize);
            }

            // Encrypt block

            let encrypted_block = encrypt(
                &self.write_buf[0..block_size],
                EncryptionMethod::Aes256Zip,
                &self.key,
            )?;

            // Remove block from buffer

            self.write_buf = self.write_buf[block_size..].to_vec();

            // Write block

            self.write_current_block(&encrypted_block)?;
        }

        if self.written_size > self.file_size {
            return Err(BlockWriteError::ExceededFileSize);
        }

        // If the file size has been reached, write the last block

        if !self.write_buf.is_empty() && self.written_size == self.file_size {
            if self.current_write_index >= self.block_count {
                return Err(BlockWriteError::ExceededFileSize);
            }

            // Encrypt block

            let encrypted_block = encrypt(&self.write_buf, EncryptionMethod::Aes256Zip, &self.key)?;

            // Clear write buf

            self.write_buf.clear();

            // Write block

            self.write_current_block(&encrypted_block)?;
        }

        Ok(())
    }

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