expanded key, helper functions

CryptoSwift/AES.swift

@@ -11,11 +11,10 @@ import Foundation
 public class AES {
     
     enum Variant:Int {
-        case aes128, aes192, aes256
+        case unknown, aes128, aes192, aes256
 
         var Nk:Int { // Nk words
-            let blocks = [4,6,8]
-            return blocks[self.rawValue]
+            return [4,6,8][self.rawValue - 1]
         }
         
         var Nb:Int { // Nb words
@@ -27,7 +26,8 @@ public class AES {
         }
     }
     
-    let variant:Variant;
+    let variant:Variant
+    let key:NSData
     
     private let sBox:[Byte] = [
         0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76, // 0
@@ -47,6 +47,31 @@ public class AES {
         0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf, // e
         0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16] // f
     
+    private let invSBox:[Byte] = [0x52, 0x09, 0x6a, 0xd5, 0x30, 0x36, 0xa5, 0x38, 0xbf, 0x40, 0xa3,
+        0x9e, 0x81, 0xf3, 0xd7, 0xfb, 0x7c, 0xe3, 0x39, 0x82, 0x9b, 0x2f,
+        0xff, 0x87, 0x34, 0x8e, 0x43, 0x44, 0xc4, 0xde, 0xe9, 0xcb, 0x54,
+        0x7b, 0x94, 0x32, 0xa6, 0xc2, 0x23, 0x3d, 0xee, 0x4c, 0x95, 0x0b,
+        0x42, 0xfa, 0xc3, 0x4e, 0x08, 0x2e, 0xa1, 0x66, 0x28, 0xd9, 0x24,
+        0xb2, 0x76, 0x5b, 0xa2, 0x49, 0x6d, 0x8b, 0xd1, 0x25, 0x72, 0xf8,
+        0xf6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xd4, 0xa4, 0x5c, 0xcc, 0x5d,
+        0x65, 0xb6, 0x92, 0x6c, 0x70, 0x48, 0x50, 0xfd, 0xed, 0xb9, 0xda,
+        0x5e, 0x15, 0x46, 0x57, 0xa7, 0x8d, 0x9d, 0x84, 0x90, 0xd8, 0xab,
+        0x00, 0x8c, 0xbc, 0xd3, 0x0a, 0xf7, 0xe4, 0x58, 0x05, 0xb8, 0xb3,
+        0x45, 0x06, 0xd0, 0x2c, 0x1e, 0x8f, 0xca, 0x3f, 0x0f, 0x02, 0xc1,
+        0xaf, 0xbd, 0x03, 0x01, 0x13, 0x8a, 0x6b, 0x3a, 0x91, 0x11, 0x41,
+        0x4f, 0x67, 0xdc, 0xea, 0x97, 0xf2, 0xcf, 0xce, 0xf0, 0xb4, 0xe6,
+        0x73, 0x96, 0xac, 0x74, 0x22, 0xe7, 0xad, 0x35, 0x85, 0xe2, 0xf9,
+        0x37, 0xe8, 0x1c, 0x75, 0xdf, 0x6e, 0x47, 0xf1, 0x1a, 0x71, 0x1d,
+        0x29, 0xc5, 0x89, 0x6f, 0xb7, 0x62, 0x0e, 0xaa, 0x18, 0xbe, 0x1b,
+        0xfc, 0x56, 0x3e, 0x4b, 0xc6, 0xd2, 0x79, 0x20, 0x9a, 0xdb, 0xc0,
+        0xfe, 0x78, 0xcd, 0x5a, 0xf4, 0x1f, 0xdd, 0xa8, 0x33, 0x88, 0x07,
+        0xc7, 0x31, 0xb1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xec, 0x5f, 0x60,
+        0x51, 0x7f, 0xa9, 0x19, 0xb5, 0x4a, 0x0d, 0x2d, 0xe5, 0x7a, 0x9f,
+        0x93, 0xc9, 0x9c, 0xef, 0xa0, 0xe0, 0x3b, 0x4d, 0xae, 0x2a, 0xf5,
+        0xb0, 0xc8, 0xeb, 0xbb, 0x3c, 0x83, 0x53, 0x99, 0x61, 0x17, 0x2b,
+        0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26, 0xe1, 0x69, 0x14, 0x63, 0x55,
+        0x21, 0x0c, 0x7d]
+    
     // Parameters for Linear Congruence Generators
     private let Rcon:[Byte] = [0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d, 0x9a,
                                0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39,
@@ -65,24 +90,34 @@ public class AES {
                                0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd,
                                0x61, 0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d]
     
-    init?(key:NSData, iv:NSData, variant:Variant = Variant.aes256) {
-        self.variant = variant
-    }
-    
-    public init() {
-        self.variant = .aes256
+    public init?(key:NSData) {
+        self.key = key
+        switch (key.length * 8) {
+        case 128:
+            self.variant = .aes128
+            break
+        case 192:
+            self.variant = .aes192
+            break
+        case 256:
+            self.variant = .aes256
+            break
+        default:
+            self.variant = .unknown
+            return nil;
+        }
     }
-    
-    func encrypt(message:NSData) -> NSData? {
         
-        return nil
+    public func encrypt(message:NSData) -> NSData? {
+        let out = cipher(message.bytes())
+        return NSData.withBytes(out)
     }
     
     func decrypt(message:NSData) -> NSData? {
         return nil
     }
     
-    public func keyExpansion(key: NSData) -> [Byte] {
+    public func expandedKey() -> [Byte] {
         
         /*
         * Function used in the Key Expansion routine that takes a four-byte
@@ -92,25 +127,21 @@ public class AES {
         func subWord(word:[Byte]) -> [Byte] {
             var result = word
             for i in 0..<4 {
-                //result[i] = sBox[Int(16*((result[i] & 0xF0) &>> 4) + (result[i] & 0x0F))]
                 result[i] = sBox[Int(word[i])]
             }
             return result
         }
 
-        //let w = NSMutableData(length: variant.Nb*(variant.Nr + 1) * 4)
-        var w:[Byte] = [Byte](count: 4 * variant.Nb * (variant.Nr + 1), repeatedValue: 0)
-        let keyBytes = key.bytes()
-        var i:Int = 0
-        while (i < variant.Nk) {
+        var w:[Byte] = [Byte](count: variant.Nb * (variant.Nr + 1) * 4, repeatedValue: 0)
+        let keyBytes:[Byte] = key.bytes()
+        for i in 0..<variant.Nk {
             w[(4*i)+0] = keyBytes[(4*i)+0]
             w[(4*i)+1] = keyBytes[(4*i)+1]
             w[(4*i)+2] = keyBytes[(4*i)+2]
             w[(4*i)+3] = keyBytes[(4*i)+3]
-            i++
         }
         
-        i = variant.Nk
+        var i = variant.Nk
         while (i < variant.Nb * (variant.Nr + 1)) {
             var tmp:[Byte] = [Byte](count: 4, repeatedValue: 0)
             tmp[0] = w[4*(i-1)+0]
@@ -136,8 +167,132 @@ public class AES {
         return w;
     }
     
-    func cipher(input:[[Byte]], output:[[Byte]], w:[[Byte]]) {
-        //var state:[[Byte]] = [[Byte]](count: 4, repeatedValue: [Byte](count: variant.Nb, repeatedValue: 0))
-        var state = input
+    private func cipher(input:[Byte]) -> [Byte] {
+        let expandedKeyW = expandedKey()
+        
+        var state:[[Byte]] = [[Byte]](count: 4, repeatedValue: [Byte](count: variant.Nb, repeatedValue: 0))
+        var idx = 0
+        for (i, row) in enumerate(state) {
+            for (j, val) in enumerate(row) {
+                state[j][i] = input[idx++]
+            }
+        }
+        
+        state = addRoundKey(state,expandedKeyW, 0)
+        
+        for roundCount in 1..<variant.Nr {
+            state = subBytes(state)
+            state = shiftRows(state)
+            state = mixColumns(state)
+            state = addRoundKey(state, expandedKeyW, roundCount)
+        }
+        
+        state = subBytes(state)
+        state = shiftRows(state)
+        state = addRoundKey(state, expandedKeyW, variant.Nr)
+        
+        var out:[Byte] = [Byte]()
+        for i in 0..<state.count {
+            for j in 0..<state[0].count {
+                out.append(state[j][i])
+            }
+        }
+        
+        return out
+    }
+}
+
+extension AES {
+    public func subBytes(state:[[Byte]]) -> [[Byte]] {
+        var result = state
+        for (i,row) in enumerate(state) {
+            for (j,value) in enumerate(row) {
+                result[i][j] = sBox[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 t:[Byte] = [Byte](count: 4, repeatedValue: 0)
+        for r in 1..<4 {
+            for c in 0..<variant.Nb {
+                t[c] = state[r][(c + r) % variant.Nb]
+            }
+            for c in 0..<variant.Nb {
+                result[r][c] = t[c]
+            }
+        }
+        return result
+    }
+    
+    public func mult(a:UInt8, _ b:UInt8) -> UInt8 {
+        var a = a, b = b
+        var p:UInt8 = 0, hbs:UInt8 = 0
+        
+        for i in 0..<8 {
+            if (b & 1 == 1) {
+                p ^= a
+            }
+            hbs = a & 0x80
+            a <<= 1
+            if (hbs > 0) {
+                a ^= 0x1B
+            }
+            b >>= 1
+        }
+        return p
+    }
+    
+    public func matrixMult(matrix:[[Byte]], _ array:[Byte]) -> [Byte] {
+        var returnArray:[Byte] = array.map({ _ in return 0 })
+        for (i, row) in enumerate(matrix) {
+            for (j, boxVal) in enumerate(row) {
+                returnArray[i] = mult(boxVal, array[j]) ^ returnArray[i]
+            }
+        }
+        return returnArray
+    }
+
+    public func addRoundKey(state:[[Byte]], _ expandedKeyW:[Byte], _ round:Int) -> [[Byte]] {
+        var newState = state.map({ val -> [Byte] in return val.map { _ in return 0 } })
+        for c in 0..<variant.Nb {
+            for i in 0..<4 {
+                newState[i][c] = state[i][c] ^ expandedKeyW[(4*variant.Nb*round)+(variant.Nb*c)+i]
+            }
+        }
+        return newState
+    }
+    
+    public func mixColumns(state:[[Byte]]) -> [[Byte]] {
+        var state = state
+        var colBox:[[Byte]] = [[2,3,1,1],
+                               [1,2,3,1],
+                               [1,1,2,3],
+                               [3,1,1,2]]
+        
+        var rowMajorState = 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 {
+                rowMajorState[j][i] = state[i][j]
+            }
+        }
+        
+        var newRowMajorState = state.map({ val -> [Byte] in return val.map { _ in return 0 } })
+        
+        for (i, row) in enumerate(rowMajorState) {
+            newRowMajorState[i] = matrixMult(colBox, row)
+        }
+        
+        for i in 0..<state.count {
+            for j in 0..<state[0].count {
+                state[i][j] = newRowMajorState[j][i]
+            }
+        }
+        
+        return state
     }
 }

CryptoSwift/Cipher.swift

@@ -10,17 +10,16 @@ import Foundation
 
 public enum Cipher {
     case ChaCha20(key: NSData, iv: NSData)
-    case AES(key: NSData, iv: NSData)
+    case AES(key: NSData)
     
     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):
+                var aes = CryptoSwift.AES(key: key)
                 return aes?.encrypt(message)
-            
         }
     }
     
@@ -29,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):
+                var aes = CryptoSwift.AES(key: key);
                 return aes?.decrypt(message)
         }
     }

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

@@ -2,6 +2,23 @@
 
 import Foundation
 
-var str:String? = ""
-var data = str?.dataUsingEncoding(NSUTF8StringEncoding, allowLossyConversion: false)
+public func mult(a:UInt8, b:UInt8) -> UInt8 {
+    var a = a, b = b
+    var p:UInt8 = 0, hbs:UInt8 = 0
+    
+    for i in 0..<8 {
+        if (b & 1 == 1) {
+            p ^= a
+        }
+        hbs = a & 0x80
+        a <<= 1
+        if (hbs != 0) {
+            a ^= 0x1B
+        }
+        println("\(i) p=\(p) a=\(a)")
+        b >>= 1
+    }
+    return p
+}
 
+mult(0x0e, 0x5f)

CryptoSwiftTests/CipherTests.swift

@@ -10,7 +10,22 @@ import Foundation
 import XCTest
 import CryptoSwift
 
+func compareMatrix(a:[[Byte]], b:[[Byte]]) -> Bool {
+    for (i,arr) in enumerate(a) {
+        for (j,val) in enumerate(arr) {
+            if (val != b[i][j]) {
+                println("\(val) vs \(b[i][j])")
+                return false
+            }
+        }
+    }
+    return true
+}
+
 class CipherTests: XCTestCase {
+    
+    // 128 bit key
+    let aesKey:[Byte] = [0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f]
 
     override func setUp() {
         super.setUp()
@@ -21,19 +36,92 @@ class CipherTests: XCTestCase {
     }
     
     func testAES() {
-        // 256 bit key
-        var key:[Byte] = [Byte](count: 32, repeatedValue: 0)
-        for i:Byte in 0..<Byte(key.count) {
-            key[Int(i)] = i
+        var input:[Byte] = [0x00, 0x11, 0x22, 0x33,
+                            0x44, 0x55, 0x66, 0x77,
+                            0x88, 0x99, 0xaa, 0xbb,
+                            0xcc, 0xdd, 0xee, 0xff];
+        
+//        if let aes = AES(key: NSData.withBytes(aesKey)) {
+//            let out = aes.encrypt(NSData.withBytes(input))
+//            println(out!)
+//        } else {
+//            XCTAssert(false, "failed")
+//        }
+    }
+    
+    func testAES_SubBytes() {
+        let input:[[Byte]] = [[0x00, 0x10, 0x20, 0x30],
+                              [0x40, 0x50, 0x60, 0x70],
+                              [0x80, 0x90, 0xa0, 0xb0],
+                              [0xc0, 0xd0, 0xe0, 0xf0]]
+        
+        let expected:[[Byte]] = [[0x63, 0xca, 0xb7, 0x04],
+                                 [0x09, 0x53, 0xd0, 0x51],
+                                 [0xcd, 0x60, 0xe0, 0xe7],
+                                 [0xba, 0x70, 0xe1, 0x8c]]
+        
+        XCTAssertTrue(compareMatrix(expected, AES(key: NSData.withBytes(aesKey))!.subBytes(input)), "subBytes failed")
+    }
+    
+    func testAES_shiftRows() {
+        let input:[[Byte]] = [[0x63, 0x09, 0xcd, 0xba],
+            [0xca, 0x53, 0x60, 0x70],
+            [0xb7, 0xd0, 0xe0, 0xe1],
+            [0x04, 0x51, 0xe7, 0x8c]]
+        
+        let expected:[[Byte]] = [[0x63, 0x9, 0xcd, 0xba],
+            [0x53, 0x60, 0x70, 0xca],
+            [0xe0, 0xe1, 0xb7, 0xd0],
+            [0x8c, 0x4, 0x51, 0xe7]]
+        
+        XCTAssertTrue(compareMatrix(expected, AES(key: NSData.withBytes(aesKey))!.shiftRows(input)), "shiftRows failed")
+    }
+    
+    func testAES_mult() {
+        XCTAssertTrue(AES(key: NSData.withBytes(aesKey))!.mult(0x0e, 0x5f) == 0x17, "Multiplication failed")
+    }
+    
+    func testAES_expandKey() {
+        let expected:[Byte] = [0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf, 0xd6, 0xaa, 0x74, 0xfd, 0xd2, 0xaf, 0x72, 0xfa, 0xda, 0xa6, 0x78, 0xf1, 0xd6, 0xab, 0x76, 0xfe, 0xb6, 0x92, 0xcf, 0xb, 0x64, 0x3d, 0xbd, 0xf1, 0xbe, 0x9b, 0xc5, 0x0, 0x68, 0x30, 0xb3, 0xfe, 0xb6, 0xff, 0x74, 0x4e, 0xd2, 0xc2, 0xc9, 0xbf, 0x6c, 0x59, 0xc, 0xbf, 0x4, 0x69, 0xbf, 0x41, 0x47, 0xf7, 0xf7, 0xbc, 0x95, 0x35, 0x3e, 0x3, 0xf9, 0x6c, 0x32, 0xbc, 0xfd, 0x5, 0x8d, 0xfd, 0x3c, 0xaa, 0xa3, 0xe8, 0xa9, 0x9f, 0x9d, 0xeb, 0x50, 0xf3, 0xaf, 0x57, 0xad, 0xf6, 0x22, 0xaa, 0x5e, 0x39, 0xf, 0x7d, 0xf7, 0xa6, 0x92, 0x96, 0xa7, 0x55, 0x3d, 0xc1, 0xa, 0xa3, 0x1f, 0x6b, 0x14, 0xf9, 0x70, 0x1a, 0xe3, 0x5f, 0xe2, 0x8c, 0x44, 0xa, 0xdf, 0x4d, 0x4e, 0xa9, 0xc0, 0x26, 0x47, 0x43, 0x87, 0x35, 0xa4, 0x1c, 0x65, 0xb9, 0xe0, 0x16, 0xba, 0xf4, 0xae, 0xbf, 0x7a, 0xd2, 0x54, 0x99, 0x32, 0xd1, 0xf0, 0x85, 0x57, 0x68, 0x10, 0x93, 0xed, 0x9c, 0xbe, 0x2c, 0x97, 0x4e, 0x13, 0x11, 0x1d, 0x7f, 0xe3, 0x94, 0x4a, 0x17, 0xf3, 0x7, 0xa7, 0x8b, 0x4d, 0x2b, 0x30, 0xc5]
+        
+        if let aes = AES(key: NSData.withBytes(aesKey)) {
+            XCTAssertEqual(expected, aes.expandedKey(), "expandKey failed")
         }
+
+    }
+    
+    func testAES_addRoundKey() {
+        let input:[[Byte]] = [[0x00, 0x44, 0x88, 0xcc],
+            [0x11, 0x55, 0x99, 0xdd],
+            [0x22, 0x66, 0xaa, 0xee],
+            [0x33, 0x77, 0xbb, 0xff]]
         
-        var input:[Byte] = [0x00, 0x11, 0x22, 0x33,
-            0x44, 0x55, 0x66, 0x77,
-            0x88, 0x99, 0xaa, 0xbb,
-            0xcc, 0xdd, 0xee, 0xff];
+        let expected:[[Byte]] = [[0, 64, 128, 192],
+            [16, 80, 144, 208],
+            [32, 96, 160, 224],
+            [48, 112, 176, 240]]
         
-        let aes = AES() // 256
-        let w = aes.keyExpansion(NSData.withBytes(key))
+        if let aes = AES(key: NSData.withBytes(aesKey)) {
+            let result = aes.addRoundKey(input, aes.expandedKey(), 0)
+            XCTAssertTrue(compareMatrix(expected, result), "addRoundKey failed")
+        }
+    }
+   
+    func testAES_mixColumns() {
+        let input:[[Byte]] = [[0x63, 0x9, 0xcd, 0xba],
+            [0x53, 0x60, 0x70, 0xca],
+            [0xe0, 0xe1, 0xb7, 0xd0],
+            [0x8c, 0x4, 0x51, 0xe7]]
+        
+        let expected:[[Byte]] = [[0x5f, 0x57, 0xf7, 0x1d],
+            [0x72, 0xf5, 0xbe, 0xb9],
+            [0x64, 0xbc, 0x3b, 0xf9],
+            [0x15, 0x92, 0x29, 0x1a]]
+        
+        if let aes = AES(key: NSData.withBytes(aesKey)) {
+            let result = aes.mixColumns(input)
+            XCTAssertTrue(compareMatrix(expected, result), "mixColumns failed")
+        }
     }
     
     func testPoly1305() {