happy_cracking/crypto/
polybius.rs

1use anyhow::Result;
2use clap::Subcommand;
3
4#[derive(Subcommand)]
5pub enum PolybiusAction {
6    #[command(about = "Encrypt with Polybius square")]
7    Encrypt {
8        #[arg(help = "Input text")]
9        input: String,
10    },
11    #[command(about = "Decrypt Polybius square")]
12    Decrypt {
13        #[arg(help = "Polybius encoded text (number pairs, space-separated)")]
14        input: String,
15    },
16}
17
18pub fn run(action: PolybiusAction) -> Result<()> {
19    match action {
20        PolybiusAction::Encrypt { input } => {
21            println!("{}", encrypt(&input)?);
22        }
23        PolybiusAction::Decrypt { input } => {
24            println!("{}", decrypt(&input)?);
25        }
26    }
27    Ok(())
28}
29
30// Default 5x5 grid: ABCDEFGHIKLMNOPQRSTUVWXYZ (I/J merged)
31const GRID: [char; 25] = [
32    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T',
33    'U', 'V', 'W', 'X', 'Y', 'Z',
34];
35
36pub fn encrypt(input: &str) -> Result<String> {
37    // Optimization: Avoid `format!` overhead in hot loop and `Vec<String>::join`
38    let mut bytes = Vec::with_capacity(input.len() * 3);
39
40    for b in input.bytes() {
41        if !b.is_ascii_alphabetic() {
42            continue;
43        }
44
45        let mut b = b.to_ascii_uppercase();
46        if b == b'J' {
47            b = b'I';
48        }
49
50        let idx = b - b'A';
51        let (row, col) = match b {
52            b'A'..=b'I' => (idx / 5 + 1, idx % 5 + 1),
53            b'K'..=b'Z' => {
54                let shifted_idx = idx - 1; // Since J is skipped
55                (shifted_idx / 5 + 1, shifted_idx % 5 + 1)
56            }
57            _ => continue,
58        };
59
60        if !bytes.is_empty() {
61            bytes.push(b' ');
62        }
63        bytes.push(b'0' + row);
64        bytes.push(b'0' + col);
65    }
66
67    if bytes.is_empty() && input.chars().any(|c| c.is_ascii_alphabetic()) {
68        anyhow::bail!("Failed to encrypt input");
69    }
70
71    Ok(unsafe { String::from_utf8_unchecked(bytes) })
72}
73
74pub fn decrypt(input: &str) -> Result<String> {
75    if input.is_empty() {
76        return Ok(String::new());
77    }
78
79    // Optimization: Avoid dynamic allocations in loop by working with bytes directly
80    let mut result = String::with_capacity(input.len() / 3);
81    for token in input.split_whitespace() {
82        let bytes = token.as_bytes();
83        if bytes.len() != 2 {
84            anyhow::bail!("Invalid Polybius pair: {}", token);
85        }
86
87        let r_byte = bytes[0];
88        let c_byte = bytes[1];
89
90        if !(b'1'..=b'5').contains(&r_byte) || !(b'1'..=b'5').contains(&c_byte) {
91            if !r_byte.is_ascii_digit() || !c_byte.is_ascii_digit() {
92                anyhow::bail!("Invalid digit in pair: {}", token);
93            }
94            anyhow::bail!("Polybius coordinates out of range (1-5): {}", token);
95        }
96
97        let row = (r_byte - b'0') as usize;
98        let col = (c_byte - b'0') as usize;
99
100        let idx = (row - 1) * 5 + (col - 1);
101        result.push(GRID[idx]);
102    }
103
104    Ok(result)
105}
happy_cracking/crypto/polybius.rs

happy_cracking/crypto/
polybius.rs
