1use crate::common::alphabet::Alphabet;
8use crate::common::cipher::Cipher;
9use crate::common::keygen::concatonated_keystream;
10use crate::common::{alphabet, substitute};
11
12pub struct Autokey {
16    key: String,
17}
18
19impl Cipher for Autokey {
20    type Key = String;
21    type Algorithm = Autokey;
22
23    fn new(key: String) -> Autokey {
30        if key.is_empty() {
31            panic!("The key must contain at least one character.");
32        } else if !alphabet::STANDARD.is_valid(&key) {
33            panic!("The key cannot contain non-alphabetic symbols.");
34        }
35
36        Autokey { key }
37    }
38
39    fn encrypt(&self, message: &str) -> Result<String, &'static str> {
52        Ok(substitute::key_substitution(
57            message,
58            &concatonated_keystream(&self.key, message),
59            |mi, ki| alphabet::STANDARD.modulo((mi + ki) as isize),
60        ))
61    }
62
63    fn decrypt(&self, ciphertext: &str) -> Result<String, &'static str> {
76        let mut plaintext = String::new();
80        let mut keystream: Vec<char> = self.key.clone().chars().collect();
81        let mut stream_idx: usize = 0;
82
83        for ct in ciphertext.chars() {
84            let ctpos = alphabet::STANDARD.find_position(ct);
85            match ctpos {
86                Some(ci) => {
87                    let decrypted_character: char;
88                    if let Some(kc) = keystream.get(stream_idx) {
89                        if let Some(ki) = alphabet::STANDARD.find_position(*kc) {
90                            let si = alphabet::STANDARD.modulo(ci as isize - ki as isize);
92                            decrypted_character =
93                                alphabet::STANDARD.get_letter(si, ct.is_uppercase());
94                        } else {
95                            panic!("Keystream contains a non-alphabetic symbol.");
96                        }
97                    } else {
98                        panic!("Keystream is not large enough for full substitution of message.");
99                    }
100
101                    plaintext.push(decrypted_character);
102                    keystream.push(decrypted_character);
103                    stream_idx += 1;
104                }
105                None => plaintext.push(ct), }
107        }
108
109        Ok(plaintext)
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[test]
118    fn with_utf8() {
119        let m = "Attack 🗡️ the east wall";
120        let a = Autokey::new(String::from("fort"));
121
122        assert_eq!(m, a.decrypt(&a.encrypt(m).unwrap()).unwrap());
123    }
124
125    #[test]
126    fn simple_encrypt_decrypt_test() {
127        let message = "defend the east wall of the castle";
128        let v = Autokey::new(String::from("fortification"));
129
130        let c_text = v.encrypt(message).unwrap();
131        let p_text = v.decrypt(&c_text).unwrap();
132
133        assert_eq!(message, p_text);
134    }
135
136    #[test]
137    fn decrypt_test() {
138        let ciphertext = "lxfopktmdcgn";
139        let v = Autokey::new(String::from("lemon"));
140        assert_eq!("attackatdawn", v.decrypt(ciphertext).unwrap());
141    }
142
143    #[test]
144    fn valid_key() {
145        Autokey::new(String::from("LeMon"));
146    }
147
148    #[test]
149    #[should_panic]
150    fn key_with_symbols() {
151        Autokey::new(String::from("!em@n"));
152    }
153
154    #[test]
155    #[should_panic]
156    fn key_with_whitespace() {
157        Autokey::new(String::from("wow this key is a real lemon"));
158    }
159}