EasyCrypto 0.8.1

use crate::Misc;


mod Aes;


pub use Aes::{Aes128Cipher, Aes128CipherKey, Aes192Cipher, Aes192CipherKey, Aes256Cipher, Aes256CipherKey};


pub trait BlockCipherTrait
{
	const BLOCK_SIZE: usize;

	fn EncryptBlock (&self, input: &[u8; Self::BLOCK_SIZE]) -> [u8; Self::BLOCK_SIZE];
	fn DecryptBlock (&self, input: &[u8; Self::BLOCK_SIZE]) -> [u8; Self::BLOCK_SIZE];

	fn EncryptBlockInplace (&self, input: &[u8; Self::BLOCK_SIZE], output: &mut [u8; Self::BLOCK_SIZE]);
	fn DecryptBlockInplace (&self, input: &[u8; Self::BLOCK_SIZE], output: &mut [u8; Self::BLOCK_SIZE]);
}


#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BlockMode
{
	ECB,
	CBC,
	CFB,
	OFB,
	CTR,
	XTS,
}


#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PaddingMode
{
	PKCS,
}


#[derive(Debug, Clone)]
pub struct BlockCipher<T: BlockCipherTrait>
where [u8; T::BLOCK_SIZE]:
{
	cipher: T,
	blockMode: BlockMode,
	paddingMode: PaddingMode,
	iv: [u8; T::BLOCK_SIZE],
}

impl<T: BlockCipherTrait> BlockCipher<T>
where [u8; T::BLOCK_SIZE]:
{
	pub fn New (cipher: T, blockMode: BlockMode, paddingMode: PaddingMode, iv: &[u8; T::BLOCK_SIZE]) -> Self
	{
		Self
		{
			cipher,
			blockMode,
			paddingMode,
			iv: iv.clone (),
		}
	}

	fn XorBlock (a: &[u8; T::BLOCK_SIZE], b: &[u8; T::BLOCK_SIZE]) -> [u8; T::BLOCK_SIZE]
	{
		let mut output = [0; T::BLOCK_SIZE];

		for i in 0..output.len ()
		{
			output[i] = a[i] ^ b[i];
		}

		output
	}

	fn XorBlockInplace (a: &[u8; T::BLOCK_SIZE], b: &mut [u8; T::BLOCK_SIZE])
	{
		for i in 0..T::BLOCK_SIZE
		{
			b[i] ^= a[i];
		}
	}

	fn CalcEncryptedSize (&self, input: &[u8]) -> usize
	{
		match self.paddingMode
		{
			PaddingMode::PKCS =>
			{
				if input.len () % T::BLOCK_SIZE == 0
				{
					input.len () + T::BLOCK_SIZE
				}
				else
				{
					Misc::Align::<usize> (input.len (), T::BLOCK_SIZE)
				}
			}
		}
	}

	fn CalcDecryptedSize (&self, output: &[u8]) -> usize
	{
		match self.paddingMode
		{
			PaddingMode::PKCS => output.len () - output[output.len () - 1] as usize,
		}
	}

	// include padding
	fn MakeLastBlock (&self, input: &[u8]) -> [u8; T::BLOCK_SIZE]
	{
		let mut output = [0; T::BLOCK_SIZE];

		match self.paddingMode
		{
			PaddingMode::PKCS =>
			{
				// calc padding size
				let leftOverSize = Misc::Align::<usize> (input.len (), T::BLOCK_SIZE) - input.len ();
				let paddingSize = T::BLOCK_SIZE - leftOverSize;

				// first part
				for i in 0..leftOverSize
				{
					output[i] = input[input.len () - leftOverSize + i];
				}

				for i in leftOverSize..T::BLOCK_SIZE
				{
					output[i] = paddingSize.try_into ().unwrap ();
				}
			}
		}

		output
	}

	// Output must big enough to hold encrypted input, included padding
	fn ECBEncryptInplace (&self, input: &[u8], output: &mut [u8])
	{
		// check output size is correct
		assert! (output.len () >= self.CalcEncryptedSize (input));
		let lastBlock = self.MakeLastBlock (input);

		for (i, inputBlock) in input.chunks_exact (T::BLOCK_SIZE).enumerate ()
		{
			let i = i*T::BLOCK_SIZE;

			self.cipher.EncryptBlockInplace (inputBlock.try_into ().unwrap (), (&mut output[i..i + T::BLOCK_SIZE]).try_into ().unwrap ());
		}

		// last block
		let outputLen = output.len ();
		self.cipher.EncryptBlockInplace (&lastBlock, (&mut output[outputLen - T::BLOCK_SIZE..outputLen]).try_into ().unwrap ());
	}

	// Output must big enough to hold encrypted input, included padding
	// Return plaintext size accroding to padding rule
	fn ECBDecryptInplace (&self, input: &[u8], output: &mut [u8]) -> usize
	{
		assert! (input.len () % T::BLOCK_SIZE == 0);
		assert! (output.len () >= input.len ());

		for (i, inputBlock) in input.chunks_exact (T::BLOCK_SIZE).enumerate ()
		{
			let i = i*T::BLOCK_SIZE;

			self.cipher.DecryptBlockInplace (inputBlock.try_into ().unwrap (), (&mut output[i..i + T::BLOCK_SIZE]).try_into ().unwrap ());
		}

		self.CalcDecryptedSize (&output[0..input.len ()])
	}

	// Output must big enough to hold encrypted input, included padding
	fn CBCEncryptInplace (&self, input: &[u8], output: &mut [u8])
	{
		// check output size is correct
		assert! (output.len () >= self.CalcEncryptedSize (input));
		let lastBlock = self.MakeLastBlock (input);

		let mut currentXorBlock = self.iv;

		for (i, inputBlock) in input.chunks_exact (T::BLOCK_SIZE).enumerate ()
		{
			let i = i*T::BLOCK_SIZE;

			Self::XorBlockInplace (inputBlock.try_into ().unwrap (), &mut currentXorBlock);
			self.cipher.EncryptBlockInplace (&currentXorBlock, (&mut output[i..i + T::BLOCK_SIZE]).try_into ().unwrap ());

			currentXorBlock = output[i..i + T::BLOCK_SIZE].try_into ().unwrap ();
		}

		// last block
		let outputLen = output.len ();
		Self::XorBlockInplace (&lastBlock, &mut currentXorBlock);
		self.cipher.EncryptBlockInplace (&currentXorBlock, (&mut output[outputLen - T::BLOCK_SIZE..outputLen]).try_into ().unwrap ());
	}

	// Output must big enough to hold decrypted input, included padding
	// Return plaintext size accroding to padding rule
	fn CBCDecryptInplace (&self, input: &[u8], output: &mut [u8]) -> usize
	{
		assert! (input.len () % T::BLOCK_SIZE == 0);
		assert! (output.len () >= input.len ());

		let mut currentXorBlock = self.iv;

		for (i, inputBlock) in input.chunks_exact (T::BLOCK_SIZE).enumerate ()
		{
			let i = i*T::BLOCK_SIZE;

			self.cipher.DecryptBlockInplace (inputBlock.try_into ().unwrap (), (&mut output[i..i + T::BLOCK_SIZE]).try_into ().unwrap ());
			Self::XorBlockInplace (&currentXorBlock, (&mut output[i..i + T::BLOCK_SIZE]).try_into ().unwrap ());

			currentXorBlock = inputBlock.try_into ().unwrap ();
		}

		self.CalcDecryptedSize (&output[0..input.len ()])
	}

	// Output must big enough to hold encrypted input, included padding
	fn OFBEncryptInplace (&self, input: &[u8], output: &mut [u8])
	{
		// check output size is correct
		assert! (output.len () >= self.CalcEncryptedSize (input));
		let lastBlock = self.MakeLastBlock (input);

		let mut tmp = self.iv;

		for (i, inputBlock) in input.chunks_exact (T::BLOCK_SIZE).enumerate ()
		{
			let i = i*T::BLOCK_SIZE;

			self.cipher.EncryptBlockInplace (&tmp, (&mut output[i..i + T::BLOCK_SIZE]).try_into ().unwrap ());
			tmp = output[i..i + T::BLOCK_SIZE].try_into ().unwrap ();
			Self::XorBlockInplace (inputBlock.try_into ().unwrap (), (&mut output[i..i + T::BLOCK_SIZE]).try_into ().unwrap ());
		}

		// last block
		let outputLen = output.len ();
		self.cipher.EncryptBlockInplace (&tmp, (&mut output[outputLen - T::BLOCK_SIZE..outputLen]).try_into ().unwrap ());
		Self::XorBlockInplace (&lastBlock, (&mut output[outputLen - T::BLOCK_SIZE..outputLen]).try_into ().unwrap ());
	}

	// Output must big enough to hold decrypted input, included padding
	// Return plaintext size accroding to padding rule
	fn OFBDecryptInplace (&self, input: &[u8], output: &mut [u8]) -> usize
	{
		assert! (input.len () % T::BLOCK_SIZE == 0);
		assert! (output.len () >= input.len ());

		// ofb decryption is exactly the same as encryption
		self.OFBEncryptInplace (input, output);

		self.CalcDecryptedSize (&output[0..input.len ()])
	}

	// Output must big enough to hold encrypted input, included padding
	fn CTREncryptInplace (&self, input: &[u8], output: &mut [u8])
	{
		fn IncreaseCounter (counter: &mut [u8])
		{
			for i in (0..counter.len ()).rev ()
			{
				if counter[i] != 0xff
				{
					counter[i] += 1;
					return;
				}
				else
				{
					counter[i] = 0;
				}
			}

			panic! ("CTR counter overflow!");
		}

		// check output size is correct
		assert! (output.len () >= self.CalcEncryptedSize (input));
		let lastBlock = self.MakeLastBlock (input);

		let mut counter = [0; T::BLOCK_SIZE];

		for (i, inputBlock) in input.chunks_exact (T::BLOCK_SIZE).enumerate ()
		{
			let i = i*T::BLOCK_SIZE;

			self.cipher.EncryptBlockInplace (&counter, (&mut output[i..i + T::BLOCK_SIZE]).try_into ().unwrap ());
			Self::XorBlockInplace (inputBlock.try_into ().unwrap (), (&mut output[i..i + T::BLOCK_SIZE]).try_into ().unwrap ());
			IncreaseCounter (&mut counter);
		}

		// last block
		let outputLen = output.len ();
		self.cipher.EncryptBlockInplace (&counter, (&mut output[outputLen - T::BLOCK_SIZE..outputLen]).try_into ().unwrap ());
		Self::XorBlockInplace (&lastBlock, (&mut output[outputLen - T::BLOCK_SIZE..outputLen]).try_into ().unwrap ());
	}

	// Output must big enough to hold decrypted input, included padding
	// Return plaintext size accroding to padding rule
	fn CTRDecryptInplace (&self, input: &[u8], output: &mut [u8]) -> usize
	{
		assert! (input.len () % T::BLOCK_SIZE == 0);
		assert! (output.len () >= input.len ());

		// ctr decryption is exactly the same as encryption
		self.CTREncryptInplace (input, output);

		self.CalcDecryptedSize (&output[0..input.len ()])
	}

	/// Output must big enough to hold encrypted input, included padding
	pub fn EncryptInplace (&self, input: &[u8], output: &mut [u8])
	{
		assert! (output.len () >= self.CalcEncryptedSize (input));

		match self.blockMode
		{
			BlockMode::ECB => self.ECBEncryptInplace (input, output),
			BlockMode::CBC => self.CBCEncryptInplace (input, output),
			BlockMode::OFB => self.OFBEncryptInplace (input, output),
			BlockMode::CTR => self.CTREncryptInplace (input, output),
			_ => panic! ("unimplemented!"),
		}
	}

	/// Output must big enough to hold decrypted input, included padding
	/// Return plaintext size accroding to padding rule
	pub fn DecryptInplace (&self, input: &[u8], output: &mut [u8]) -> usize
	{
		assert! (input.len () % T::BLOCK_SIZE == 0);
		assert! (output.len () >= input.len ());

		match self.blockMode
		{
			BlockMode::ECB => self.ECBDecryptInplace (input, output),
			BlockMode::CBC => self.CBCDecryptInplace (input, output),
			BlockMode::OFB => self.OFBDecryptInplace (input, output),
			BlockMode::CTR => self.CTRDecryptInplace (input, output),
			_ => panic! ("unimplemented!"),
		}
	}

	pub fn Encrypt (&self, input: &[u8]) -> Vec<u8>
	{
		let mut output = vec! [0; self.CalcEncryptedSize (input)];
		self.EncryptInplace (input, &mut output);

		output
	}

	// Decrypt the input, trip the output according to padding mode or not
	pub fn Decrypt (&self, input: &[u8], stripOutput: bool) -> Vec<u8>
	{
		let mut output = vec! [0; self.CalcEncryptedSize (input)];
		let plainTextLen = self.DecryptInplace (input, &mut output);

		if stripOutput
		{
			output.resize (plainTextLen, 0);
		}

		output
	}
}


#[cfg(test)]
mod tests
{
	use super::*;

	#[test]
	fn Aes128ECB ()
	{
		let key = [1; 16];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes128Cipher = Aes128Cipher::New (&Aes128CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes128Cipher, BlockMode::ECB, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes128CBC ()
	{
		let key = [1; 16];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes128Cipher = Aes128Cipher::New (&Aes128CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes128Cipher, BlockMode::CBC, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes128OFB ()
	{
		let key = [1; 16];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes128Cipher = Aes128Cipher::New (&Aes128CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes128Cipher, BlockMode::OFB, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes128CTR ()
	{
		let key = [1; 16];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes128Cipher = Aes128Cipher::New (&Aes128CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes128Cipher, BlockMode::CTR, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes192ECB ()
	{
		let key = [1; 24];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes192Cipher = Aes192Cipher::New (&Aes192CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes192Cipher, BlockMode::ECB, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes192CBC ()
	{
		let key = [1; 24];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes192Cipher = Aes192Cipher::New (&Aes192CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes192Cipher, BlockMode::CBC, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes192OFB ()
	{
		let key = [1; 24];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes192Cipher = Aes192Cipher::New (&Aes192CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes192Cipher, BlockMode::OFB, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes192CTR ()
	{
		let key = [1; 24];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes192Cipher = Aes192Cipher::New (&Aes192CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes192Cipher, BlockMode::CTR, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes256ECB ()
	{
		let key = [1; 32];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes256Cipher = Aes256Cipher::New (&Aes256CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes256Cipher, BlockMode::ECB, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes256CBC ()
	{
		let key = [1; 32];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes256Cipher = Aes256Cipher::New (&Aes256CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes256Cipher, BlockMode::CBC, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes256OFB ()
	{
		let key = [1; 32];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes256Cipher = Aes256Cipher::New (&Aes256CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes256Cipher, BlockMode::OFB, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}

	#[test]
	fn Aes256CTR ()
	{
		let key = [1; 32];
		let iv = [2; 16];
		let plainText = [3; 16];

		let aes256Cipher = Aes256Cipher::New (&Aes256CipherKey::New (&key));
		let blockCipher = BlockCipher::New (aes256Cipher, BlockMode::CTR, PaddingMode::PKCS, &iv);

		let cipherText = blockCipher.Encrypt (&plainText);
		let plainText2 = blockCipher.Decrypt (&cipherText, true);

		assert! (plainText2.len () == 16);
		assert! (plainText == <[u8; 16]>::try_from (plainText2).unwrap ());
	}
}