use crate::crc::crc32_byte;
use std::io::{self, Read, Write};
pub(crate) const HEADER_LEN: usize = 12;
#[derive(Debug, Clone)]
pub(crate) struct Cipher {
keys: [u32; 3],
}
impl Cipher {
pub(crate) fn new(password: &[u8]) -> Self {
let mut cipher = Cipher {
keys: [0x1234_5678, 0x2345_6789, 0x3456_7890],
};
for &byte in password {
cipher.update(byte);
}
cipher
}
#[inline]
fn update(&mut self, byte: u8) {
self.keys[0] = crc32_byte(self.keys[0], byte);
self.keys[1] = self.keys[1]
.wrapping_add(self.keys[0] & 0xFF)
.wrapping_mul(134_775_813)
.wrapping_add(1);
self.keys[2] = crc32_byte(self.keys[2], (self.keys[1] >> 24) as u8);
}
#[inline]
fn keystream_byte(&self) -> u8 {
let temp = (self.keys[2] | 2) as u16;
(temp.wrapping_mul(temp ^ 1) >> 8) as u8
}
#[inline]
pub(crate) fn decrypt(&mut self, buf: &mut [u8]) {
for byte in buf {
let plain = *byte ^ self.keystream_byte();
self.update(plain);
*byte = plain;
}
}
#[inline]
pub(crate) fn encrypt(&mut self, buf: &mut [u8]) {
for byte in buf {
let cipher = *byte ^ self.keystream_byte();
self.update(*byte);
*byte = cipher;
}
}
}
#[derive(Debug)]
pub struct Decryptor<R> {
reader: R,
cipher: Cipher,
check_byte: u8,
}
impl<R: Read> Decryptor<R> {
pub fn new(mut reader: R, password: &[u8]) -> io::Result<Self> {
let mut cipher = Cipher::new(password);
let mut header = [0u8; HEADER_LEN];
reader.read_exact(&mut header)?;
cipher.decrypt(&mut header);
Ok(Decryptor {
reader,
cipher,
check_byte: header[HEADER_LEN - 1],
})
}
pub fn check_byte(&self) -> u8 {
self.check_byte
}
pub fn into_inner(self) -> R {
self.reader
}
}
impl<R: Read> Read for Decryptor<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if buf.is_empty() {
return Ok(0);
}
let read = self.reader.read(buf)?;
self.cipher.decrypt(&mut buf[..read]);
Ok(read)
}
}
const CHUNK: usize = 8 * 1024;
#[derive(Debug)]
pub struct Encryptor<W> {
writer: W,
cipher: Cipher,
scratch: Box<[u8; CHUNK]>,
}
impl<W: Write> Encryptor<W> {
pub fn new(
mut writer: W,
password: &[u8],
header_random: [u8; 11],
check_byte: u8,
) -> io::Result<Self> {
let mut cipher = Cipher::new(password);
let mut header = [0u8; HEADER_LEN];
header[..HEADER_LEN - 1].copy_from_slice(&header_random);
header[HEADER_LEN - 1] = check_byte;
cipher.encrypt(&mut header);
writer.write_all(&header)?;
Ok(Encryptor {
writer,
cipher,
scratch: Box::new([0u8; CHUNK]),
})
}
pub fn into_inner(self) -> W {
self.writer
}
pub fn get_ref(&self) -> &W {
&self.writer
}
pub fn get_mut(&mut self) -> &mut W {
&mut self.writer
}
}
impl<W: Write> Write for Encryptor<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
for chunk in buf.chunks(CHUNK) {
let part = &mut self.scratch[..chunk.len()];
part.copy_from_slice(chunk);
self.cipher.encrypt(part);
self.writer.write_all(part)?;
}
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
self.writer.flush()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn encrypt_decrypt_round_trip() {
let plaintext = b"the quick brown fox";
let mut buf = *plaintext;
Cipher::new(b"hunter2").encrypt(&mut buf);
assert_ne!(&buf, plaintext);
Cipher::new(b"hunter2").decrypt(&mut buf);
assert_eq!(&buf, plaintext);
}
#[test]
fn encryptor_decryptor_round_trip() {
let password = b"hunter2";
let check_byte = 0x42;
let plaintext = b"the quick brown fox jumps over the lazy dog".repeat(400);
let mut encrypted = Vec::new();
let mut encryptor =
Encryptor::new(&mut encrypted, password, [7u8; HEADER_LEN - 1], check_byte).unwrap();
encryptor.write_all(&plaintext).unwrap();
encryptor.flush().unwrap();
assert_eq!(encrypted.len(), HEADER_LEN + plaintext.len());
let mut decryptor = Decryptor::new(encrypted.as_slice(), password).unwrap();
assert_eq!(decryptor.check_byte(), check_byte);
let mut decrypted = Vec::new();
decryptor.read_to_end(&mut decrypted).unwrap();
assert_eq!(decrypted, plaintext);
}
}