Padding protocol as parameter for AES cipher

CryptoSwift.xcodeproj/project.pbxproj

@@ -18,6 +18,7 @@
 		7547195119931802002FA5F1 /* IntExtension.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7547195019931802002FA5F1 /* IntExtension.swift */; };
 		754BE45B19693E190098E6F3 /* CryptoSwift.h in Headers */ = {isa = PBXBuildFile; fileRef = 754BE45A19693E190098E6F3 /* CryptoSwift.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		754BE46819693E190098E6F3 /* HashTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 754BE46719693E190098E6F3 /* HashTests.swift */; };
+		754C30B71AA13BC000E6FFA4 /* PKCS7.swift in Sources */ = {isa = PBXBuildFile; fileRef = 754C30B61AA13BC000E6FFA4 /* PKCS7.swift */; };
 		754C8FED19979F94005AD904 /* ArrayExtension.swift in Sources */ = {isa = PBXBuildFile; fileRef = 754C8FEC19979F94005AD904 /* ArrayExtension.swift */; };
 		754DD76E19A149AF00E52288 /* CryptoHashBase.swift in Sources */ = {isa = PBXBuildFile; fileRef = 754DD76D19A149AF00E52288 /* CryptoHashBase.swift */; };
 		755111E819B7B7DF00C2AD86 /* Authenticator.swift in Sources */ = {isa = PBXBuildFile; fileRef = 755111E719B7B7DF00C2AD86 /* Authenticator.swift */; };
@@ -129,6 +130,7 @@
 		754BE46019693E190098E6F3 /* CryptoSwiftTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = CryptoSwiftTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
 		754BE46619693E190098E6F3 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
 		754BE46719693E190098E6F3 /* HashTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = HashTests.swift; sourceTree = "<group>"; };
+		754C30B61AA13BC000E6FFA4 /* PKCS7.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = PKCS7.swift; sourceTree = "<group>"; };
 		754C8FEC19979F94005AD904 /* ArrayExtension.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ArrayExtension.swift; sourceTree = "<group>"; };
 		754DD76D19A149AF00E52288 /* CryptoHashBase.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = CryptoHashBase.swift; sourceTree = "<group>"; };
 		755111E719B7B7DF00C2AD86 /* Authenticator.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Authenticator.swift; sourceTree = "<group>"; };
@@ -215,6 +217,8 @@
 				7563B2E719B14D4300B152CD /* Cipher.swift */,
 				75BC3AE21A4E412000ADF343 /* CipherBlockMode.swift */,
 				757DA2501A4E5E95002BA3EF /* PaddingMode.swift */,
+				75A663A51AA0CAD00052110B /* Padding.swift */,
+				754C30B61AA13BC000E6FFA4 /* PKCS7.swift */,
 				755111E719B7B7DF00C2AD86 /* Authenticator.swift */,
 				754DD76D19A149AF00E52288 /* CryptoHashBase.swift */,
 				750A545F1992D2680017DA75 /* MD5.swift */,
@@ -225,7 +229,6 @@
 				751C5C3C19B26B000094C75D /* Poly1305.swift */,
 				758A94251A65AEB100E46135 /* HMAC.swift */,
 				75A74B261A1FF6B2004419F1 /* AES.swift */,
-				75A663A51AA0CAD00052110B /* Padding.swift */,
 				759D481019B517BC005FF7FC /* BitExtension.swift */,
 				758F3F771992F6CE0014BBDA /* UInt8Extension.swift */,
 				7547195019931802002FA5F1 /* IntExtension.swift */,
@@ -415,6 +418,7 @@
 				757EF7F519AAA82400586276 /* CRC.swift in Sources */,
 				75D94E2619B60C4F007CB2A4 /* UInt32Extension.swift in Sources */,
 				75A74B271A1FF6B2004419F1 /* AES.swift in Sources */,
+				754C30B71AA13BC000E6FFA4 /* PKCS7.swift in Sources */,
 				755111E819B7B7DF00C2AD86 /* Authenticator.swift in Sources */,
 				754DD76E19A149AF00E52288 /* CryptoHashBase.swift in Sources */,
 				758C764119B61AE500653BC6 /* Generics.swift in Sources */,

CryptoSwift/AES.swift

@@ -140,14 +140,15 @@ public class AES {
     
     :returns: Encrypted data
     */
-    public func encrypt(message:NSData, addPadding: Bool = true) -> NSData? {
+
+    public func encrypt(message:NSData, padding: Padding? = PKCS7()) -> NSData? {
         var finalMessage = message;
 
-        if (addPadding) {
-            finalMessage = PKCS7.add(message, blockSize: AES.blockSizeBytes())
+        if let padding = padding {
+            finalMessage = padding.add(message, blockSize: AES.blockSizeBytes())
         } else if (message.length % AES.blockSizeBytes() != 0) {
             // 128 bit block exceeded, need padding
-            assertionFailure("AES 128-bit block exceeded!")
+            assert(false, "AES 128-bit block exceeded!");
             return nil
         }
         
@@ -206,7 +207,7 @@ public class AES {
         }
         
         if (out != nil && removePadding) {
-            return PKCS7.remove(NSData.withBytes(out!))
+            return PKCS7().remove(NSData.withBytes(out!))
         }
         
         return out == nil ? nil : NSData.withBytes(out!)

CryptoSwift/Cipher.swift

@@ -43,7 +43,7 @@ public enum Cipher {
                 return chacha?.encrypt(message)
             case .AES(let key, let iv, let blockMode):
                 var aes = CryptoSwift.AES(key: key, iv: iv, blockMode: blockMode)
-                return aes?.encrypt(message)
+                return aes?.encrypt(message, padding: PKCS7())
         }
     }
     

CryptoSwift/PKCS7.swift

@@ -0,0 +1,44 @@
+//
+//  PKCS7.swift
+//  CryptoSwift
+//
+//  Created by Marcin Krzyzanowski on 28/02/15.
+//  Copyright (c) 2015 Marcin Krzyzanowski. All rights reserved.
+//
+
+import Foundation
+
+public struct PKCS7: Padding {
+    
+    public init() {
+        
+    }
+    
+    public func add(data: NSData , blockSize:Int) -> NSData {
+        var padding = UInt8(blockSize) - (UInt8(data.length) % UInt8(blockSize))
+        var withPadding = NSMutableData(data: data)
+        if (padding == 0) {
+            // If the original data is a multiple of N bytes, then an extra block of bytes with value N is added.
+            for i in 0..<blockSize {
+                withPadding.appendBytes([UInt8(blockSize)])
+            }
+        } else {
+            // The value of each added byte is the number of bytes that are added
+            for i in 0..<padding {
+                withPadding.appendBytes([padding])
+            }
+        }
+        return withPadding
+    }
+    
+    public func remove(data: NSData, blockSize:Int? = nil) -> NSData
+    {
+        var padding:UInt8 = 0
+        data.subdataWithRange(NSRange(location: data.length - 1, length: 1)).getBytes(&padding, length: 1)
+        
+        if padding >= 1 {
+            return data.subdataWithRange(NSRange(location: 0, length: data.length - Int(padding)))
+        }
+        return data
+    }
+}

CryptoSwift/Padding.swift

@@ -9,37 +9,6 @@
 import Foundation
 
 public protocol Padding {
-    static func add(data: NSData, blockSize:Int) -> NSData;
-    static func remove(data: NSData, blockSize:Int?) -> NSData;
-}
-
-public struct PKCS7: Padding {
-    
-    public static func add(data: NSData , blockSize:Int) -> NSData {
-        var padding = UInt8(blockSize) - (UInt8(data.length) % UInt8(blockSize))
-        var withPadding = NSMutableData(data: data)
-        if (padding == 0) {
-            // If the original data is a multiple of N bytes, then an extra block of bytes with value N is added.
-            for i in 0..<blockSize {
-                withPadding.appendBytes([UInt8(blockSize)])
-            }
-        } else {
-            // The value of each added byte is the number of bytes that are added
-            for i in 0..<padding {
-                withPadding.appendBytes([padding])
-            }
-        }
-        return withPadding
-    }
-    
-    public static func remove(data: NSData, blockSize:Int? = nil) -> NSData
-    {
-        var padding:UInt8 = 0
-        data.subdataWithRange(NSRange(location: data.length - 1, length: 1)).getBytes(&padding, length: 1)
-        
-        if padding >= 1 {
-            return data.subdataWithRange(NSRange(location: 0, length: data.length - Int(padding)))
-        }
-        return data
-    }
+    func add(data: NSData, blockSize:Int) -> NSData;
+    func remove(data: NSData, blockSize:Int?) -> NSData;
 }

CryptoSwift/Playground/CryptoPlayground.playground/section-1.swift

@@ -7,7 +7,27 @@ let plaintext = "Lorem ipsum"
 let MD5 = plaintext.md5()
 
 
-func foo(param:String) -> Bool {
-    assertionFailure("Don't")
-    return false
-}
+protocol Padding {
+    func foo() -> String
+}
+
+struct PKCS7: Padding {
+    func foo() -> String {
+        return "dupa"
+    }
+}
+
+func createTest<T:Padding>(p: T) -> T{
+    return p
+}
+
+func test<T:Padding>(p: T) -> String {
+    return p.foo()
+}
+
+func testtest() {
+    let t = createTest(PKCS7())
+    println(t.foo())
+}
+
+testtest()

CryptoSwiftTests/AESTests.swift

@@ -22,7 +22,7 @@ class AESTests: XCTestCase {
         let expected:[UInt8] = [0xae,0x8c,0x59,0x95,0xb2,0x6f,0x8e,0x3d,0xb0,0x6f,0x0a,0xa5,0xfe,0xc4,0xf0,0xc2];
         
         if let aes = AES(key: NSData.withBytes(key), iv: NSData.withBytes(iv), blockMode: .CBC) {
-            let encrypted = aes.encrypt(NSData.withBytes(input), addPadding: false)
+            let encrypted = aes.encrypt(NSData.withBytes(input), padding: nil)
             XCTAssertEqual(encrypted!, NSData.withBytes(expected), "encryption failed")
             let decrypted = aes.decrypt(encrypted!)
             XCTAssertEqual(decrypted!, NSData.withBytes(input), "decryption failed")
@@ -43,7 +43,7 @@ class AESTests: XCTestCase {
             0x70, 0xb4, 0xc5, 0x5a];
         
         if let aes = AES(key: NSData.withBytes(aesKey), blockMode: .ECB) {
-            let encrypted = aes.encrypt(NSData.withBytes(input), addPadding: false)
+            let encrypted = aes.encrypt(NSData.withBytes(input), padding: nil)
             XCTAssertEqual(encrypted!, NSData.withBytes(expected), "encryption failed")
             let decrypted = aes.decrypt(encrypted!, removePadding: false)
             XCTAssertEqual(decrypted!, NSData.withBytes(input), "decryption failed")
@@ -60,7 +60,7 @@ class AESTests: XCTestCase {
         
         if let aes = AES(key: NSData.withBytes(key), iv:NSData.withBytes(iv), blockMode: .CBC) {
             XCTAssertTrue(aes.blockMode == .CBC, "Invalid block mode")
-            let encrypted = aes.encrypt(NSData.withBytes(plaintext), addPadding: false)
+            let encrypted = aes.encrypt(NSData.withBytes(plaintext), padding: nil)
             XCTAssertEqual(encrypted!, NSData.withBytes(expected), "encryption failed")
             let decrypted = aes.decrypt(encrypted!, removePadding: false)
             XCTAssertEqual(decrypted!, NSData.withBytes(plaintext), "decryption failed")
@@ -77,7 +77,7 @@ class AESTests: XCTestCase {
         
         if let aes = AES(key: NSData.withBytes(key), iv:NSData.withBytes(iv), blockMode: .CFB) {
             XCTAssertTrue(aes.blockMode == .CFB, "Invalid block mode")
-            let encrypted = aes.encrypt(NSData.withBytes(plaintext), addPadding: false)
+            let encrypted = aes.encrypt(NSData.withBytes(plaintext), padding: nil)
             XCTAssertEqual(encrypted!, NSData.withBytes(expected), "encryption failed")
             let decrypted = aes.decrypt(encrypted!, removePadding: false)
             XCTAssertEqual(decrypted!, NSData.withBytes(plaintext), "decryption failed")

CryptoSwiftTests/PaddingTests.swift

@@ -14,27 +14,27 @@ class PaddingTests: XCTestCase {
     func testPKCS7_0() {
         let input:[UInt8]    = [1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6]
         let expected:[UInt8] = [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.add(NSData.withBytes(input), blockSize: 16)
+        let padded = PKCS7().add(NSData.withBytes(input), blockSize: 16)
         XCTAssertEqual(padded, NSData.withBytes(expected), "PKCS7 failed")
-        let clean = PKCS7.remove(padded)
+        let clean = PKCS7().remove(padded)
         XCTAssertEqual(clean, NSData.withBytes(input), "PKCS7 failed")
     }
     
     func testPKCS7_1() {
         let input:[UInt8]    = [1,2,3,4,5,6,7,8,9,0,1,2,3,4,5]
         let expected:[UInt8] = [1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,1]
-        let padded = PKCS7.add(NSData.withBytes(input), blockSize: 16)
+        let padded = PKCS7().add(NSData.withBytes(input), blockSize: 16)
         XCTAssertEqual(padded, NSData.withBytes(expected), "PKCS7 failed")
-        let clean = PKCS7.remove(padded)
+        let clean = PKCS7().remove(padded)
         XCTAssertEqual(clean, NSData.withBytes(input), "PKCS7 failed")
     }
     
     func testPKCS7_2() {
         let input:[UInt8]    = [1,2,3,4,5,6,7,8,9,0,1,2,3,4]
         let expected:[UInt8] = [1,2,3,4,5,6,7,8,9,0,1,2,3,4,2,2]
-        let padded = PKCS7.add(NSData.withBytes(input), blockSize: 16)
+        let padded = PKCS7().add(NSData.withBytes(input), blockSize: 16)
         XCTAssertEqual(padded, NSData.withBytes(expected), "PKCS7 failed")
-        let clean = PKCS7.remove(padded)
+        let clean = PKCS7().remove(padded)
         XCTAssertEqual(clean, NSData.withBytes(input), "PKCS7 failed")
     }
 }