encrypt decrypt test

src/main.rs

@@ -2,6 +2,7 @@ mod lfsr;
 mod tea3;
 
 use lfsr::Lfsr;
+use tea3::Tea3;
 
 fn main() {
     let mut lfsr = Lfsr::new(4, vec![0, 3], vec![0x12, 0x34, 0x56, 0x78]);
@@ -10,4 +11,41 @@ fn main() {
         let byte = lfsr.next();
         println!("{:02x} | state: {:?}", byte, lfsr.state());
     }
+
+    let key = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+    let state = vec![0; 8];
+
+    let mut cipher = Tea3::new(key.clone(), state.clone());
+    cipher.init();
+
+    println!("\nKeystream:");
+    for _ in 0..10 {
+        let byte = cipher.next_byte();
+        println!("{:#04x}", byte);
+    }
+
+    let plaintext = b"hello world";
+
+    let mut cipher_enc = Tea3::new(key.clone(), state.clone());
+    cipher_enc.init();
+
+    let ciphertext: Vec<u8> = plaintext
+        .iter()
+        .map(|&b| b ^ cipher_enc.next_byte())
+        .collect();
+
+    println!("\nPlaintext : {:?}", plaintext);
+    println!("Ciphertext: {:?}", ciphertext);
+
+    let mut cipher_dec = Tea3::new(key, state);
+    cipher_dec.init();
+
+    let decrypted: Vec<u8> = ciphertext
+        .iter()
+        .map(|&b| b ^ cipher_dec.next_byte())
+        .collect();
+
+    println!("Decrypted : {:?}", decrypted);
+
+    assert_eq!(plaintext.to_vec(), decrypted);
 }

src/tea3.rs

@@ -8,6 +8,9 @@ pub struct Tea3 {
     state_register: Vec<u8>,
 }
 
+/// tea-3 implementation
+/// new + init to create the registers and init the tea-3 construction
+/// next_byte to get the next cypher byte
 impl Tea3 {
     pub fn new(key_register: Vec<u8>, state_register: Vec<u8>) -> Self {
         assert_eq!(key_register.len(), 10, "TEA3 key register is 10B");
@@ -66,15 +69,23 @@ impl Tea3 {
         r7
     }
 
+    pub fn step(&mut self) -> u8 {
+        let kout: u8 = self.step_key_register();
+        self.step_state_register(kout)
+    }
+
+    pub fn init(&mut self) {
+        for _i in 0..32 {
+            self.step();
+        }
+    }
+
     /// outputs the next cipher byte
     pub fn next_byte(&mut self) -> u8 {
-        let mut kout: u8;
-        for _i in 1..19 {
-            kout = self.step_key_register();
-            self.step_state_register(kout);
+        for _i in 0..18 {
+            self.step();
         }
-        kout = self.step_key_register();
-        self.step_state_register(kout)
+        self.step()
     }
 
     pub fn key_register(&self) -> &[u8] {
@@ -352,4 +363,35 @@ mod tests {
         assert_ne!(s[0], 1);
         assert_ne!(s[5], 5);
     }
+
+    #[test]
+    fn test_encrypt_decrypt() {
+        // Example key (10B) and state (8B)
+        let key = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+        // should be derived from CC, CN and LA, placeholder value
+        let state = vec![0; 8];
+
+        let mut cipher = Tea3::new(key.clone(), state.clone());
+        cipher.init();
+
+        let plaintext = b"hello world";
+
+        let mut cipher_enc = Tea3::new(key.clone(), state.clone());
+        cipher_enc.init();
+
+        let ciphertext: Vec<u8> = plaintext
+            .iter()
+            .map(|&b| b ^ cipher_enc.next_byte())
+            .collect();
+
+        let mut cipher_dec = Tea3::new(key, state);
+        cipher_dec.init();
+
+        let decrypted: Vec<u8> = ciphertext
+            .iter()
+            .map(|&b| b ^ cipher_dec.next_byte())
+            .collect();
+
+        assert_eq!(plaintext.to_vec(), decrypted);
+    }
 }