problem 9

3 files changed, 19 insertions, 0 deletions

diff --git a/src/lib.rs b/src/lib.rs
index fcde2ce..028faa7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -11,6 +11,7 @@ pub use aes::decrypt_aes_128_ecb;
 pub use aes::find_aes_128_ecb_encrypted_string;
 pub use base64::to_base64;
 pub use primitives::fixed_xor;
+pub use primitives::pad_pkcs7;
 pub use primitives::repeating_key_xor;
 pub use xor::find_single_byte_xor_encrypted_string;
 pub use xor::crack_single_byte_xor;
diff --git a/src/primitives.rs b/src/primitives.rs
index 0551e27..b0d7116 100644
--- a/src/primitives.rs
+++ b/src/primitives.rs
@@ -1,3 +1,5 @@
+use std;
+
 pub fn fixed_xor (bytes1: &[u8], bytes2: &[u8]) -> Vec<u8> {
     return bytes1.iter()
         .zip(bytes2.iter())
@@ -21,6 +23,15 @@ pub fn hamming (bytes1: &[u8], bytes2: &[u8]) -> u64 {
     count_bits(&fixed_xor(bytes1, bytes2)[..])
 }
 
+pub fn pad_pkcs7 (block: &[u8], blocksize: u8) -> Vec<u8> {
+    let padding_bytes = blocksize - (block.len() % blocksize as usize) as u8;
+    return block
+        .iter()
+        .map(|c| *c)
+        .chain(std::iter::repeat(padding_bytes).take(padding_bytes as usize))
+        .collect();
+}
+
 fn count_bits (bytes: &[u8]) -> u64 {
     bytes.iter().map(|&c| { count_bits_byte(c) }).fold(0, |acc, n| acc + n)
 }
diff --git a/tests/lib.rs b/tests/lib.rs
index a3e797e..e5392b0 100644
--- a/tests/lib.rs
+++ b/tests/lib.rs
@@ -111,3 +111,10 @@ fn problem_8 () {
     let got = matasano::find_aes_128_ecb_encrypted_string(&possibles[..]);
     assert_eq!(got, ciphertext);
 }
+
+#[test]
+fn problem_9 () {
+    let block = b"YELLOW SUBMARINE";
+    let got = matasano::pad_pkcs7(block, 20);
+    assert_eq!(got, b"YELLOW SUBMARINE\x04\x04\x04\x04");
+}