From 67953c47d914a3b6af40b61cff337e8134cb0186 Mon Sep 17 00:00:00 2001
From: Jixun Wu <8041017-jixunmoe@users.noreply.gitlab.com>
Date: Wed, 11 Sep 2024 02:02:16 +0100
Subject: [PATCH] fix: fix wrapper method

---
 src/lib.rs | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index e25d27f..f68eed4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@
 //!
 //! Notably, it uses a different round number and uses a "tweaked" CBC mode.
 
+use crate::cbc::get_encrypted_size;
 use byteorder::{ByteOrder, BE};
 use thiserror::Error;
 
@@ -56,9 +57,10 @@ pub fn parse_key(key: &[u8]) -> Result<[u32; 4], TcTeaError> {
 pub fn encrypt<T: AsRef<[u8]>>(plaintext: T, key: &[u8]) -> Result<Vec<u8>, TcTeaError> {
     let key = parse_key(key)?;
     let plaintext = plaintext.as_ref();
-    let mut cipher = vec![0u8; plaintext.len()];
-    let result = cbc::decrypt(&mut cipher, plaintext, &key)?;
-    Ok(Vec::from(result))
+    let cipher_len = get_encrypted_size(plaintext.len());
+    let mut cipher = vec![0u8; cipher_len];
+    cbc::encrypt(&mut cipher, plaintext, &key)?;
+    Ok(cipher)
 }
 
 /// Decrypts a byte array containing the following:
@@ -77,3 +79,14 @@ pub fn decrypt<T: AsRef<[u8]>>(encrypted: T, key: &[u8]) -> Result<Vec<u8>, TcTe
     let result = cbc::decrypt(&mut plain, encrypted, &key)?;
     Ok(Vec::from(result))
 }
+
+#[test]
+fn test_sanity_test() -> Result<(), TcTeaError> {
+    let key = b"43218765dcbahgfe";
+    let message = b"this is a test message.";
+    let cipher = encrypt(message, key)?;
+    let plain = decrypt(cipher, key)?;
+    assert_eq!(message, plain.as_slice());
+
+    Ok(())
+}