Skip to main content

happy_cracking/crypto/
columnar.rs

1use anyhow::Result;
2use clap::Subcommand;
3
4use super::shared::column_order;
5
6#[derive(Subcommand)]
7pub enum ColumnarAction {
8    #[command(about = "Encrypt with Columnar Transposition cipher")]
9    Encrypt {
10        #[arg(help = "Input text")]
11        input: String,
12        #[arg(short, long, help = "Keyword for column ordering")]
13        key: String,
14    },
15    #[command(about = "Decrypt Columnar Transposition cipher")]
16    Decrypt {
17        #[arg(help = "Encrypted text")]
18        input: String,
19        #[arg(short, long, help = "Keyword for column ordering")]
20        key: String,
21    },
22}
23
24pub fn run(action: ColumnarAction) -> Result<()> {
25    match action {
26        ColumnarAction::Encrypt { input, key } => {
27            println!("{}", encrypt(&input, &key)?);
28        }
29        ColumnarAction::Decrypt { input, key } => {
30            println!("{}", decrypt(&input, &key)?);
31        }
32    }
33    Ok(())
34}
35
36// Safety limit to prevent massive memory allocation (DoS).
37// If a user provides a tiny input but a massive key (e.g. 100 million chars),
38// the padding logic would allocate 100 million 'X's.
39const MAX_KEY_LEN: usize = 1_000_000;
40
41pub fn encrypt(input: &str, key: &str) -> Result<String> {
42    if key.len() > MAX_KEY_LEN {
43        anyhow::bail!(
44            "Key exceeds maximum length of {} to prevent Denial of Service",
45            MAX_KEY_LEN
46        );
47    }
48
49    if key.is_empty() || !key.chars().all(|c| c.is_ascii_alphabetic()) {
50        anyhow::bail!("Key must be non-empty and contain only alphabetic characters");
51    }
52
53    if input.is_empty() {
54        return Ok(String::new());
55    }
56
57    let key_len = key.len();
58    let order = column_order(key);
59
60    // Pad input with 'X' to fill the grid
61    let mut padded: Vec<char> = input.chars().collect();
62    while !padded.len().is_multiple_of(key_len) {
63        padded.push('X');
64    }
65    let num_rows = padded.len() / key_len;
66
67    // Invert `order` into rank -> original column index.
68    let mut col_at_rank = vec![0usize; key_len];
69    for (col, &rank) in order.iter().enumerate() {
70        col_at_rank[rank] = col;
71    }
72
73    // Read columns in key order by indexing directly into the padded buffer.
74    // Skipping the Vec<Vec<char>> grid removes `total_len` char copies and
75    // `key_len + 1` heap allocations per encrypt call.
76    let mut result = String::with_capacity(padded.len());
77    for &col in &col_at_rank {
78        for row in 0..num_rows {
79            result.push(padded[row * key_len + col]);
80        }
81    }
82
83    Ok(result)
84}
85
86pub fn decrypt(input: &str, key: &str) -> Result<String> {
87    if key.len() > MAX_KEY_LEN {
88        anyhow::bail!(
89            "Key exceeds maximum length of {} to prevent Denial of Service",
90            MAX_KEY_LEN
91        );
92    }
93
94    if key.is_empty() || !key.chars().all(|c| c.is_ascii_alphabetic()) {
95        anyhow::bail!("Key must be non-empty and contain only alphabetic characters");
96    }
97
98    if input.is_empty() {
99        return Ok(String::new());
100    }
101
102    let key_len = key.len();
103    let chars: Vec<char> = input.chars().collect();
104    let total_len = chars.len();
105
106    if !total_len.is_multiple_of(key_len) {
107        anyhow::bail!("Ciphertext length must be a multiple of key length");
108    }
109
110    let num_rows = total_len / key_len;
111    let order = column_order(key);
112
113    // `order[col]` is the rank at which column `col` appears in the ciphertext,
114    // so its data occupies `chars[order[col] * num_rows .. (order[col]+1) * num_rows]`.
115    // Indexing straight into `chars` avoids the `Vec<Vec<char>>` column buffer
116    // and saves `total_len` char copies plus `key_len + 1` allocations.
117    let mut result = String::with_capacity(total_len);
118    for row in 0..num_rows {
119        for col in 0..key_len {
120            result.push(chars[order[col] * num_rows + row]);
121        }
122    }
123
124    Ok(result)
125}