use crate::{input, Cipher, CipherInputError, CipherResult};
pub struct FourSquare {
key1: String,
key2: String,
alphabet: String,
pad: u8,
}
impl FourSquare {
pub fn new(key1: &str, key2: &str, alphabet: &str, pad: char) -> Self {
assert_eq!(alphabet.len(), 25, "alphabet` must be 25 chars in length");
input::is_ascii(alphabet).expect("`alphabet` must be valid ascii");
input::no_repeated_chars(alphabet).expect("`alphabet` cannot contain repeated chars");
assert_eq!(key1.len(), 25, "`key1` must be 25 chars in length");
input::no_repeated_chars(key1).expect("`key1` cannot contain repeated chars");
input::in_alphabet(key1, alphabet)
.expect("all chars in `key1` must be contained in `alphabet`");
assert_eq!(key2.len(), 25, "`key2` must be 25 chars in length");
input::no_repeated_chars(key2).expect("`key2` cannot contain repeated chars");
input::in_alphabet(key2, alphabet)
.expect("all chars in `key2` must be contained in `alphabet`");
assert_ne!(alphabet.find(pad), None, "`alphabet` must contain `pad`");
Self {
key1: String::from(key1),
key2: String::from(key2),
alphabet: String::from(alphabet),
pad: pad as u8,
}
}
}
impl Cipher for FourSquare {
fn encipher(&self, ptext: &str) -> CipherResult {
input::in_alphabet(ptext, &self.alphabet)?;
let mut ptext: Vec<u8> = ptext.bytes().collect();
if ptext.len() % 2 != 0 {
ptext.push(self.pad);
}
let key1 = self.key1.as_bytes();
let key2 = self.key2.as_bytes();
let mut ctext = Vec::with_capacity(ptext.len());
for i in (0..ptext.len()).step_by(2) {
let yx1 = match self.alphabet.bytes().position(|c| c == ptext[i]) {
Some(val) => val,
None => return Err(CipherInputError::NotInAlphabet),
};
let yx2 = match self.alphabet.bytes().position(|c| c == ptext[i + 1]) {
Some(val) => val,
None => return Err(CipherInputError::NotInAlphabet),
};
let (y1, x1) = (yx1 / 5, yx1 % 5);
let (y2, x2) = (yx2 / 5, yx2 % 5);
ctext.push(key1[y1 * 5 + x2]);
ctext.push(key2[y2 * 5 + x1]);
}
Ok(String::from_utf8(ctext).unwrap())
}
fn decipher(&self, ctext: &str) -> CipherResult {
input::in_alphabet(ctext, &self.alphabet)?;
if ctext.len() % 2 != 0 {
return Err(CipherInputError::BadInput(String::from(
"`ctext` must contain an even number of chars",
)));
}
let ctext = ctext.as_bytes();
let mut ptext = Vec::with_capacity(ctext.len());
for i in (0..ctext.len()).step_by(2) {
let yx1 = match self.key1.find(|c| c == ctext[i] as char) {
Some(val) => val,
None => return Err(CipherInputError::NotInAlphabet),
};
let yx2 = match self.key2.find(|c| c == ctext[i + 1] as char) {
Some(val) => val,
None => return Err(CipherInputError::NotInAlphabet),
};
let (y1, x2) = (yx1 / 5, yx1 % 5);
let (y2, x1) = (yx2 / 5, yx2 % 5);
ptext.push(self.alphabet.as_bytes()[y1 * 5 + x1]);
ptext.push(self.alphabet.as_bytes()[y2 * 5 + x2]);
}
Ok(String::from_utf8(ptext).unwrap())
}
}