remove PKCS7 padding

krzyzanowskim · Dec 28, 2014 · d80417c · d80417c
1 parent 4b49441
commit d80417c
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 9 deletions.
diff --git a/CryptoSwift/Cipher.swift b/CryptoSwift/Cipher.swift
@@ -10,15 +10,15 @@ import Foundation
 
 public enum Cipher {
     case ChaCha20(key: NSData, iv: NSData)
-    case AES(key: NSData, iv: NSData)
+    case AES(key: NSData, iv: NSData, blockMode: CipherBlockMode)
 
     public func encrypt(message: NSData) -> NSData? {
         switch (self) {
             case .ChaCha20(let key, let iv):
                 var chacha = CryptoSwift.ChaCha20(key: key, iv: iv)
                 return chacha?.encrypt(message)
-            case .AES(let key, let iv):
-                var aes = CryptoSwift.AES(key: key, iv: iv)
+            case .AES(let key, let iv, let blockMode):
+                var aes = CryptoSwift.AES(key: key, iv: iv, blockMode: blockMode)
                 return aes?.encrypt(message)
         }
     }
@@ -28,8 +28,8 @@ public enum Cipher {
             case .ChaCha20(let key, let iv):
                 var chacha = CryptoSwift.ChaCha20(key: key, iv: iv);
                 return chacha?.decrypt(message)
-            case .AES(let key, let iv):
-                var aes = CryptoSwift.AES(key: key, iv: iv);
+            case .AES(let key, let iv, let blockMode):
+                var aes = CryptoSwift.AES(key: key, iv: iv, blockMode: blockMode);
                 return aes?.decrypt(message)
         }
     }

diff --git a/CryptoSwift/PKCS7.swift b/CryptoSwift/PKCS7.swift
@@ -31,4 +31,11 @@ public struct PKCS7 {
         }
         return withPadding
     }
+
+    public func removePadding() -> NSData {
+        if let padding = data.bytes().last {
+            return data.subdataWithRange(NSRange(location: 0, length: data.length - Int(padding)))
+        }
+        return data
+    }
 }
diff --git a/CryptoSwift/PaddingMode.swift b/CryptoSwift/PaddingMode.swift
@@ -19,4 +19,14 @@ public enum PaddingMode:Int {
             return data
         }
     }
+
+    public func removePadding(data:NSData) -> NSData {
+        switch (self) {
+        case PKCS_7:
+            return PKCS7(data: data).removePadding()
+        case None:
+            return data
+        }
+    }
+
 }
diff --git a/CryptoSwiftTests/PaddingTests.swift b/CryptoSwiftTests/PaddingTests.swift
@@ -11,24 +11,30 @@ import XCTest
 import CryptoSwift
 
 class PaddingTests: XCTestCase {
-    func testPKCS7_1() {
+    func testPKCS7_0() {
         let input:[Byte]    = [1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6]
         let expected:[Byte] = [1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16]
         let padded = PKCS7(data: NSData.withBytes(input)).addPadding(16)
         XCTAssertEqual(padded, NSData.withBytes(expected), "PKCS7 failed")
+        let clean = PKCS7(data: padded).removePadding()
+        XCTAssertEqual(clean, NSData.withBytes(input), "PKCS7 failed")
     }
 
-    func testPKCS7_2() {
+    func testPKCS7_1() {
         let input:[Byte]    = [1,2,3,4,5,6,7,8,9,0,1,2,3,4,5]
         let expected:[Byte] = [1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,1]
         let padded = PKCS7(data: NSData.withBytes(input)).addPadding(16)
         XCTAssertEqual(padded, NSData.withBytes(expected), "PKCS7 failed")
+        let clean = PKCS7(data: padded).removePadding()
+        XCTAssertEqual(clean, NSData.withBytes(input), "PKCS7 failed")
     }
 
-    func testPKCS7_3() {
+    func testPKCS7_2() {
         let input:[Byte]    = [1,2,3,4,5,6,7,8,9,0,1,2,3,4]
         let expected:[Byte] = [1,2,3,4,5,6,7,8,9,0,1,2,3,4,2,2]
         let padded = PKCS7(data: NSData.withBytes(input)).addPadding(16)
         XCTAssertEqual(padded, NSData.withBytes(expected), "PKCS7 failed")
+        let clean = PKCS7(data: padded).removePadding()
+        XCTAssertEqual(clean, NSData.withBytes(input), "PKCS7 failed")
     }
 }
diff --git a/README.md b/README.md
@@ -98,7 +98,9 @@ Working with Ciphers
 	// DECRYPT
 	let decryptedChaCha20 = Cipher.ChaCha20(setup).decrypt(encryptedData)
 	let decryptedAES = Cipher.AES(setup).decrypt(encryptedData)
-	// remember to remove padding if applied
+
+	// remove padding IF applied on encryption
+	let decryptedRaw = PKCS7(data: decryptedAES).removePadding()
 
 
 using extensions