AES invert operation methods

CryptoSwift/AES.swift

@@ -232,16 +232,32 @@ extension AES {
         return result
     }
     
+    public func invSubBytes(state:[[Byte]]) -> [[Byte]] {
+        var result = state
+        for (i,row) in enumerate(state) {
+            for (j,value) in enumerate(row) {
+                result[i][j] = invSBox[Int(value)]
+            }
+        }
+        return result
+    }
+    
     // Applies a cyclic shift to the last 3 rows of a state matrix.
     public func shiftRows(state:[[Byte]]) -> [[Byte]] {
         var result = state
-        var tmp:[Byte] = [Byte](count: 4, repeatedValue: 0)
         for r in 1..<4 {
             for c in 0..<variant.Nb {
-                tmp[c] = state[r][(c + r) % variant.Nb]
+                result[r][c] = state[r][(c + r) % variant.Nb]
             }
+        }
+        return result
+    }
+    
+    public func invShiftRows(state:[[Byte]]) -> [[Byte]] {
+        var result = state
+        for r in 1..<4 {
             for c in 0..<variant.Nb {
-                result[r][c] = tmp[c]
+                result[r][(c + r) % variant.Nb] = state[r][c]
             }
         }
         return result
@@ -313,4 +329,31 @@ extension AES {
         
         return state
     }
+    
+    public func invMixColumns(state:[[Byte]]) -> [[Byte]] {
+        var state = state
+        var invColBox:[[Byte]] = [[14,11,13,9],[9,14,11,13],[13,9,14,11],[11,13,9,14]]
+        
+        var colOrderState = state.map({ val -> [Byte] in return val.map { _ in return 0 } }) // zeroing
+        
+        for i in 0..<state.count {
+            for j in 0..<state[0].count {
+                colOrderState[j][i] = state[i][j]
+            }
+        }
+        
+        var newState = state.map({ val -> [Byte] in return val.map { _ in return 0 } })
+        
+        for (i, row) in enumerate(colOrderState) {
+            newState[i] = matrixMultiplyPolys(invColBox, row)
+        }
+        
+        for i in 0..<state.count {
+            for j in 0..<state[0].count {
+                state[i][j] = newState[j][i]
+            }
+        }
+        
+        return state
+    }
 }

CryptoSwiftTests/AESTests.swift

@@ -74,7 +74,10 @@ class AESTests: XCTestCase {
             [0xcd, 0x60, 0xe0, 0xe7],
             [0xba, 0x70, 0xe1, 0x8c]]
         
-        XCTAssertTrue(compareMatrix(expected, AES(key: NSData.withBytes(aesKey))!.subBytes(input)), "subBytes failed")
+        let substituted = AES(key: NSData.withBytes(aesKey))!.subBytes(input)
+        XCTAssertTrue(compareMatrix(expected, substituted), "subBytes failed")
+        let inverted = AES(key: NSData.withBytes(aesKey))!.invSubBytes(substituted)
+        XCTAssertTrue(compareMatrix(input, inverted), "invSubBytes failed")
     }
     
     func testAES_shiftRows() {
@@ -88,7 +91,10 @@ class AESTests: XCTestCase {
             [0xe0, 0xe1, 0xb7, 0xd0],
             [0x8c, 0x4, 0x51, 0xe7]]
         
-        XCTAssertTrue(compareMatrix(expected, AES(key: NSData.withBytes(aesKey))!.shiftRows(input)), "shiftRows failed")
+        let shifted = AES(key: NSData.withBytes(aesKey))!.shiftRows(input)
+        XCTAssertTrue(compareMatrix(expected, shifted), "shiftRows failed")
+        let inverted = AES(key: NSData.withBytes(aesKey))!.invShiftRows(shifted)
+        XCTAssertTrue(compareMatrix(input, inverted), "invShiftRows failed")
     }
     
     func testAES_multiply() {
@@ -133,8 +139,10 @@ class AESTests: XCTestCase {
             [0x15, 0x92, 0x29, 0x1a]]
         
         if let aes = AES(key: NSData.withBytes(aesKey)) {
-            let result = aes.mixColumns(input)
-            XCTAssertTrue(compareMatrix(expected, result), "mixColumns failed")
+            let mixed = aes.mixColumns(input)
+            XCTAssertTrue(compareMatrix(expected, mixed), "mixColumns failed")
+            let inverted = aes.invMixColumns(mixed)
+            XCTAssertTrue(compareMatrix(input, inverted), "invMixColumns failed")
         }
     }