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    let pad_byte = *data.last().unwrap();
82    let pad_len = pad_byte as usize;
83
84    if pad_len == 0 || pad_len > block_size {
85        anyhow::bail!("Invalid PKCS7 padding value: {}", pad_byte);
86    }
87    if data.len() < pad_len {
88        anyhow::bail!("Padding length exceeds data length");
89    }
90    if !data[data.len() - pad_len..].iter().all(|&b| b == pad_byte) {
91        anyhow::bail!("Invalid PKCS7 padding bytes");
92    }
93
94    Ok(data[..data.len() - pad_len].to_vec())
95}
96
97// Safety limit for zero-pad block size to prevent Denial of Service.
98// Without a limit, a user could pass e.g. block_size=4294967295 with a tiny
99// input, causing allocation of ~4 GB of zero bytes (OOM / memory exhaustion).
100// 16 MB is generous for any realistic block-cipher block size.
101const MAX_ZERO_PAD_BLOCK_SIZE: usize = 16 * 1024 * 1024;
102
103pub fn zero_pad(data: &[u8], block_size: usize) -> Result<Vec<u8>> {
104    if block_size == 0 || data.is_empty() {
105        return Ok(data.to_vec());
106    }
107    if block_size > MAX_ZERO_PAD_BLOCK_SIZE {
108        anyhow::bail!(
109            "Block size {} exceeds maximum of {} to prevent excessive memory allocation",
110            block_size,
111            MAX_ZERO_PAD_BLOCK_SIZE
112        );
113    }
114    let remainder = data.len() % block_size;
115    if remainder == 0 {
116        return Ok(data.to_vec());
117    }
118    let pad_len = block_size - remainder;
119    let mut result = data.to_vec();
120    result.extend(std::iter::repeat_n(0u8, pad_len));
121    Ok(result)
122}
123
124pub fn zero_unpad(data: &[u8]) -> Vec<u8> {
125    let end = data.iter().rposition(|&b| b != 0).map_or(0, |pos| pos + 1);
126    data[..end].to_vec()
127}