pmv_encryption_rs 1.0.0

// Encryption

use std::{fmt::Display, io::Write};

use aes::{
    Aes256,
    cipher::{BlockEncryptMut, KeyIvInit, block_padding::NoPadding},
};
use byteorder::{BigEndian, ByteOrder};
use cbc::Encryptor;
use flate2::{Compression, write::ZlibEncoder};
use rand::{RngCore, SeedableRng, rngs::StdRng};

use crate::EncryptionMethod;

type Aes256Cbc = Encryptor<Aes256>;

/// Encryption error
#[derive(Debug, Clone)]
pub enum EncryptionError {
    /// Invalid key provided
    InvalidKey,
    /// Internal error
    InternalError {
        /// The error message
        message: String,
    },
}

impl Display for EncryptionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            EncryptionError::InvalidKey => write!(f, "Invalid encryption key provided"),
            EncryptionError::InternalError { message } => write!(f, "{}", message),
        }
    }
}

/// Encrypts data
///
/// Parameters:
///
///  - `data` - The data to encrypt
///  - `method` - The encryption method
///  - `key` - The encryption key
///
/// Returns the encrypted bytes, or an error
pub fn encrypt(
    data: &[u8],
    method: EncryptionMethod,
    key: &[u8],
) -> Result<Vec<u8>, EncryptionError> {
    if data.is_empty() {
        return Ok(vec![0; 0]);
    }

    match method {
        EncryptionMethod::Aes256Zip => encrypt_aes_cbc_zip(data, key),
        EncryptionMethod::Aes256Flat => encrypt_aes_cbc_flat(data, key),
    }
}

const AES_BLOCK_SIZE: usize = 16;
const AES_IV_SIZE: usize = 16;

fn encrypt_aes_cbc_flat(data: &[u8], key: &[u8]) -> Result<Vec<u8>, EncryptionError> {
    if key.len() != 32 {
        return Err(EncryptionError::InvalidKey);
    }

    // Encryption method mark

    let mut method_mark: Vec<u8> = vec![0; 2];

    BigEndian::write_u16(&mut method_mark, EncryptionMethod::Aes256Flat.to_u16());

    // Header

    let mut header: Vec<u8> = vec![0; 20];

    // Include pre-encryption size to the header
    BigEndian::write_u32(&mut header[0..4], data.len() as u32);

    // Pad data
    let mut final_data = add_padding(data, AES_BLOCK_SIZE);

    // Generate IV
    let iv = generate_aes_iv();

    // Include IV into the header
    header[4..4 + AES_IV_SIZE].copy_from_slice(&iv);

    // Create cypher

    let cypher = Aes256Cbc::new(key.into(), iv.as_slice().into());

    // Encrypt data

    let final_data_len = final_data.len();
    if let Err(e) = cypher.encrypt_padded_mut::<NoPadding>(&mut final_data, final_data_len) {
        return Err(EncryptionError::InternalError {
            message: e.to_string(),
        });
    }

    Ok([method_mark, header, final_data].concat())
}

fn encrypt_aes_cbc_zip(data: &[u8], key: &[u8]) -> Result<Vec<u8>, EncryptionError> {
    if key.len() != 32 {
        return Err(EncryptionError::InvalidKey);
    }

    // Encryption method mark

    let mut method_mark: Vec<u8> = vec![0; 2];

    BigEndian::write_u16(&mut method_mark, EncryptionMethod::Aes256Zip.to_u16());

    // Compress the data

    let mut compressor = ZlibEncoder::new(Vec::new(), Compression::default());

    if let Err(e) = compressor.write_all(data) {
        return Err(EncryptionError::InternalError {
            message: e.to_string(),
        });
    }

    let compressed_bytes = match compressor.finish() {
        Ok(b) => b,
        Err(e) => {
            return Err(EncryptionError::InternalError {
                message: e.to_string(),
            });
        }
    };

    // Header

    let mut header: Vec<u8> = vec![0; 20];

    // Include pre-encryption size to the header
    BigEndian::write_u32(&mut header[0..4], compressed_bytes.len() as u32);

    // Pad data
    let mut final_data = add_padding(&compressed_bytes, AES_BLOCK_SIZE);

    // Generate IV
    let iv = generate_aes_iv();

    // Include IV into the header
    header[4..4 + AES_IV_SIZE].copy_from_slice(&iv);

    // Create cypher

    let cypher = Aes256Cbc::new(key.into(), iv.as_slice().into());

    // Encrypt data

    let final_data_len = final_data.len();
    if let Err(e) = cypher.encrypt_padded_mut::<NoPadding>(&mut final_data, final_data_len) {
        return Err(EncryptionError::InternalError {
            message: e.to_string(),
        });
    }

    Ok([method_mark, header, final_data].concat())
}

fn generate_aes_iv() -> Vec<u8> {
    let mut iv: Vec<u8> = vec![0; AES_IV_SIZE];
    let mut rng = StdRng::from_os_rng();
    rng.fill_bytes(&mut iv);

    iv
}

fn add_padding(data: &[u8], block_size: usize) -> Vec<u8> {
    let padding_size = block_size - (data.len() % block_size);
    let padding_bytes: Vec<u8> = vec![0; padding_size];

    [data, &padding_bytes].concat()
}