From 336ea83612435a9dd28dd71d96ec35fe53c335f7 Mon Sep 17 00:00:00 2001
From: lionarius <main@lionarius.ru>
Date: Tue, 3 Sep 2024 15:51:06 +0300
Subject: [PATCH] polybius

---
 src/main.rs            |  9 +++---
 src/polybius_square.rs | 73 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 4 deletions(-)
 create mode 100644 src/polybius_square.rs

diff --git a/src/main.rs b/src/main.rs
index 8dd1d3e..db8704f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,14 +1,15 @@
 mod atbash;
 mod caesar;
+mod polybius_square;
 mod util;
 
 fn main() -> anyhow::Result<()> {
-    let alphabet = "абвгдежзиклмнопрстуфхцчшщъыьэюя"
+    let alphabet = "пцукезщшгнхъф.ылоравд,жэюмитьб_счя!"
         .chars()
         .collect::<Vec<_>>();
-    let encoder = caesar::CaesarCipher::new(alphabet, 1)?;
-    println!("{}", encoder.encode("абвгд")?);
-    println!("{}", encoder.decode("бвгде")?);
+    let encoder = polybius_square::PolybiusSquareCipher::new(alphabet, 5, 7)?;
+    println!("{}", encoder.encode("таина")?);
+    println!("{}", encoder.decode("чэсыэ")?);
 
     Ok(())
 }
diff --git a/src/polybius_square.rs b/src/polybius_square.rs
new file mode 100644
index 0000000..431e396
--- /dev/null
+++ b/src/polybius_square.rs
@@ -0,0 +1,73 @@
+use crate::util;
+
+pub struct PolybiusSquareCipher {
+    alphabet: Vec<char>,
+    rows: usize,
+    columns: usize,
+}
+
+impl PolybiusSquareCipher {
+    pub fn new(
+        alphabet: impl Into<Vec<char>>,
+        columns: usize,
+        rows: usize,
+    ) -> anyhow::Result<Self> {
+        let alphabet = alphabet.into();
+        util::verify_alphabet(&alphabet)?;
+        if rows * columns != alphabet.len() {
+            return Err(anyhow::anyhow!(
+                "rows * columns must equal the length of the alphabet"
+            ));
+        }
+
+        Ok(Self {
+            alphabet,
+            rows,
+            columns,
+        })
+    }
+
+    pub fn encode(&self, input: &str) -> anyhow::Result<String> {
+        let mut output = String::with_capacity(input.len());
+
+        for c in input.chars() {
+            let index = self
+                .alphabet
+                .iter()
+                .position(|&x| x == c)
+                .ok_or(anyhow::anyhow!("invalid character"))?;
+
+            let row = index / self.columns;
+            let col = index % self.columns;
+
+            output.push(self.alphabet[((row + 1) % self.rows) * self.columns + col]);
+        }
+
+        Ok(output)
+    }
+
+    pub fn decode(&self, input: &str) -> anyhow::Result<String> {
+        let mut output = String::with_capacity(input.len());
+
+        for c in input.chars() {
+            let index = self
+                .alphabet
+                .iter()
+                .position(|&x| x == c)
+                .ok_or(anyhow::anyhow!("invalid character"))?;
+
+            let mut row = index / self.columns;
+            let col = index % self.columns;
+
+            if row == 0 {
+                row = self.rows - 1;
+            } else {
+                row -= 1;
+            }
+
+            output.push(self.alphabet[(row % self.rows) * self.columns + col]);
+        }
+
+        Ok(output)
+    }
+}