Skip to main content

happy_cracking/crypto/
padding.rs

1use anyhow::{Context, Result};
2use clap::Subcommand;
3
4#[derive(Subcommand)]
5pub enum PaddingAction {
6    #[command(about = "Apply PKCS7 padding")]
7    Pkcs7Pad {
8        #[arg(help = "Input hex string")]
9        input: String,
10        #[arg(long, default_value = "16", help = "Block size (1-255)")]
11        block_size: usize,
12    },
13    #[command(about = "Remove PKCS7 padding")]
14    Pkcs7Unpad {
15        #[arg(help = "Input hex string")]
16        input: String,
17        #[arg(long, default_value = "16", help = "Block size (1-255)")]
18        block_size: usize,
19    },
20    #[command(about = "Apply zero padding")]
21    ZeroPad {
22        #[arg(help = "Input hex string")]
23        input: String,
24        #[arg(long, default_value = "16", help = "Block size")]
25        block_size: usize,
26    },
27    #[command(about = "Remove zero padding")]
28    ZeroUnpad {
29        #[arg(help = "Input hex string")]
30        input: String,
31    },
32}
33
34pub fn run(action: PaddingAction) -> Result<()> {
35    match action {
36        PaddingAction::Pkcs7Pad { input, block_size } => {
37            let data = hex::decode(input.trim()).context("Failed to decode input hex")?;
38            let padded = pkcs7_pad(&data, block_size)?;
39            println!("{}", hex::encode(&padded));
40        }
41        PaddingAction::Pkcs7Unpad { input, block_size } => {
42            let data = hex::decode(input.trim()).context("Failed to decode input hex")?;
43            let unpadded = pkcs7_unpad(&data, block_size)?;
44            println!("{}", hex::encode(&unpadded));
45        }
46        PaddingAction::ZeroPad { input, block_size } => {
47            let data = hex::decode(input.trim()).context("Failed to decode input hex")?;
48            let padded = zero_pad(&data, block_size)?;
49            println!("{}", hex::encode(&padded));
50        }
51        PaddingAction::ZeroUnpad { input } => {
52            let data = hex::decode(input.trim()).context("Failed to decode input hex")?;
53            let unpadded = zero_unpad(&data);
54            println!("{}", hex::encode(&unpadded));
55        }
56    }
57    Ok(())
58}
59
60pub fn pkcs7_pad(data: &[u8], block_size: usize) -> Result<Vec<u8>> {
61    if block_size == 0 || block_size > 255 {
62        anyhow::bail!("Block size must be between 1 and 255");
63    }
64    let pad_len = block_size - (data.len() % block_size);
65    let mut result = data.to_vec();
66    result.extend(std::iter::repeat_n(pad_len as u8, pad_len));
67    Ok(result)
68}
69
70pub fn pkcs7_unpad(data: &[u8], block_size: usize) -> Result<Vec<u8>> {
71    if block_size == 0 || block_size > 255 {
72        anyhow::bail!("Block size must be between 1 and 255");
73    }
74    if data.is_empty() {
75        anyhow::bail!("Input data is empty");
76    }
77    if !data.len().is_multiple_of(block_size) {
78        anyhow::bail!("Input length is not a multiple of block size");
79    }
80
81    // Defense-in-depth: use ok_or_else instead of unwrap to avoid a panic
82    // if the empty-check guard above is ever refactored away.
83    let pad_byte = *data
84        .last()
85        .ok_or_else(|| anyhow::anyhow!("Input data is empty"))?;
86    let pad_len = pad_byte as usize;
87
88    if pad_len == 0 || pad_len > block_size {
89        anyhow::bail!("Invalid PKCS7 padding value: {}", pad_byte);
90    }
91    if data.len() < pad_len {
92        anyhow::bail!("Padding length exceeds data length");
93    }
94    if !data[data.len() - pad_len..].iter().all(|&b| b == pad_byte) {
95        anyhow::bail!("Invalid PKCS7 padding bytes");
96    }
97
98    Ok(data[..data.len() - pad_len].to_vec())
99}
100
101// Safety limit for zero-pad block size to prevent Denial of Service.
102// Without a limit, a user could pass e.g. block_size=4294967295 with a tiny
103// input, causing allocation of ~4 GB of zero bytes (OOM / memory exhaustion).
104// 16 MB is generous for any realistic block-cipher block size.
105const MAX_ZERO_PAD_BLOCK_SIZE: usize = 16 * 1024 * 1024;
106
107pub fn zero_pad(data: &[u8], block_size: usize) -> Result<Vec<u8>> {
108    if block_size == 0 || data.is_empty() {
109        return Ok(data.to_vec());
110    }
111    if block_size > MAX_ZERO_PAD_BLOCK_SIZE {
112        anyhow::bail!(
113            "Block size {} exceeds maximum of {} to prevent excessive memory allocation",
114            block_size,
115            MAX_ZERO_PAD_BLOCK_SIZE
116        );
117    }
118    let remainder = data.len() % block_size;
119    if remainder == 0 {
120        return Ok(data.to_vec());
121    }
122    let pad_len = block_size - remainder;
123    let mut result = data.to_vec();
124    result.extend(std::iter::repeat_n(0u8, pad_len));
125    Ok(result)
126}
127
128pub fn zero_unpad(data: &[u8]) -> Vec<u8> {
129    let end = data.iter().rposition(|&b| b != 0).map_or(0, |pos| pos + 1);
130    data[..end].to_vec()
131}