use crate::compression::{EnkryptitCompress, EnkryptitDecompress};
use crate::encryption::encryption_primitives::{decrypt_chunk, encrypt_chunk};
use crate::errors::EnkryptitError;
use crate::types::CHUNK_SIZE;
use crate::types::CompressionType;
use chacha20poly1305::KeyInit;
use chacha20poly1305::XChaCha20Poly1305;
use gradient_bar::GradientProgressBar;
use std::io::{Read, Write};
pub fn encrypt_stream<R: Read, W: Write>(
writer: &mut W,
mut reader: R,
master_nonce: [u8; 24],
key: &[u8; 32],
compression: CompressionType,
progress_bar: Option<GradientProgressBar>,
) -> Result<u64, EnkryptitError> {
let mut step: u64 = 0;
let mut buffer = vec![0u8; CHUNK_SIZE];
let mut bytes_written: u64 = 0;
let mut n = reader.read(&mut buffer)?;
let mut total_processed: u64 = 0;
let cipher = XChaCha20Poly1305::new(key.into());
let mut output = vec![0u8; CHUNK_SIZE];
while n > 0 {
let mut next_buffer = vec![0u8; CHUNK_SIZE];
let next_n = reader.read(&mut next_buffer)?;
let data = buffer[..n].to_vec();
data.compress(&mut output, compression)?;
encrypt_chunk(&mut output, &master_nonce, &cipher, step)?;
let chunk_len = output.len() as u32;
writer.write_all(&chunk_len.to_le_bytes())?;
writer.write_all(&output)?;
bytes_written += 4 + output.len() as u64;
total_processed += n as u64;
if let Some(pb) = &progress_bar {
pb.update(total_processed);
}
buffer = next_buffer;
n = next_n;
step += 1;
}
writer.write_all(b"ENK1END")?;
bytes_written += 7;
if let Some(pb) = &progress_bar {
pb.finish();
}
Ok(bytes_written)
}
pub fn decrypt_stream<R: Read, W: Write>(
writer: &mut W,
mut reader: R,
key: &[u8; 32],
compression: CompressionType,
master_nonce: [u8; 24],
progress_bar: Option<GradientProgressBar>,
) -> Result<u64, EnkryptitError> {
let mut step: u64 = 0;
let mut bytes_consumed: u64 = 0;
let mut total_processed: u64 = 0;
let cipher = XChaCha20Poly1305::new(key.into());
let mut output = vec![0u8; CHUNK_SIZE];
loop {
let mut len_buf = [0u8; 4];
match reader.read_exact(&mut len_buf) {
Ok(_) => {}
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
break;
}
Err(e) => return Err(e.into()),
}
bytes_consumed += 4;
if &len_buf == b"ENK1" {
let mut end = [0u8; 3];
reader.read_exact(&mut end)?;
if &end != b"END" {
}
break;
}
let len = u32::from_le_bytes(len_buf) as usize;
let mut payload = vec![0u8; len];
reader.read_exact(&mut payload)?;
bytes_consumed += len as u64;
decrypt_chunk(&mut payload, &cipher, &master_nonce, step)?;
payload.decompress(&mut output, compression)?;
total_processed += len as u64;
if let Some(pb) = &progress_bar {
pb.update(total_processed);
}
writer.write_all(&output)?;
step += 1;
}
if let Some(pb) = &progress_bar {
pb.finish();
}
Ok(bytes_consumed)
}