Key generation

Sources/CryptoSwift/RSA.swift

@@ -65,6 +65,25 @@ public final class RSA {
     }
   }
   
+  /// Initialize with a generated key pair
+  /// - Parameter keySize: The size of the modulus
+  public convenience init(keySize: Int) {
+    // Generate prime numbers
+    let p = BigUInt.generatePrime(keySize / 2)
+    let q = BigUInt.generatePrime(keySize / 2)
+    
+    // Calculate modulus
+    let n = p * q
+    
+    // Calculate public and private exponent
+    let e: BigUInt = 65537
+    let phi = (p - 1) * (q - 1)
+    let d = e.inverse(phi)
+    
+    // Initialize
+    self.init(n: n, e: e, d: d)
+  }
+  
   // TODO: Add initializer from PEM (ASN.1 with DER header)
   
   // TODO: Add export to PEM (ASN.1 with DER header)
@@ -93,3 +112,20 @@ extension RSA: Cipher {
   }
   
 }
+
+// MARK: BigUInt extension
+
+extension BigUInt {
+  
+  public static func generatePrime(_ width: Int) -> BigUInt {
+    // Note: Need to find a better way to generate prime numbers
+    while true {
+      var random = BigUInt.randomInteger(withExactWidth: width)
+      random |= BigUInt(1)
+      if random.isPrime() {
+        return random
+      }
+    }
+  }
+  
+}

Tests/CryptoSwiftTests/RSATests.swift

@@ -95,6 +95,22 @@ final class RSATests: XCTestCase {
     let decrypted = try! rsa.decrypt(encrypted)
     XCTAssertEqual(decrypted, message, "decrypt failed")
   }
+  
+  func testGenerateKeyPair() {
+    /*
+     * To test key generation and its validity
+     */
+    let message: Array<UInt8> = [
+      0x11, 0x22, 0x33, 0x44
+    ]
+    
+    let rsa = RSA(keySize: 2048)
+    // Sometimes the modulus size is 2047 bits, but it's okay (with two 1024 bits primes)
+    //XCTAssertEqual(rsa.keySize, 2048, "key size is not correct")
+    
+    let decrypted = try! rsa.decrypt(try! rsa.encrypt(message))
+    XCTAssertEqual(decrypted, message, "encrypt+decrypt failed")
+  }
 
 }
 
@@ -102,7 +118,8 @@ extension RSATests {
   static func allTests() -> [(String, (RSATests) -> () -> Void)] {
     let tests = [
       ("testSmallRSA", testSmallRSA),
-      ("testRSA1", testRSA1)
+      ("testRSA1", testRSA1),
+      ("testGenerateKeyPair", testGenerateKeyPair)
     ]
 
     return tests