From d2bcdfb8382f2021a01b953798153a41093c61ef Mon Sep 17 00:00:00 2001
From: Kevin Wooten <kevin@outfoxx.io>
Date: Mon, 12 Jun 2023 22:30:44 -0700
Subject: [PATCH] Add unencrypted PKCS#8 export and import

---
 Sources/ShieldSecurity/SecKeyPair.swift | 43 +++++++++++++++++++------
 Tests/SecKeyPairTests.swift             | 26 +++++++++++++--
 2 files changed, 57 insertions(+), 12 deletions(-)

diff --git a/Sources/ShieldSecurity/SecKeyPair.swift b/Sources/ShieldSecurity/SecKeyPair.swift
index 88bd5bd2f..1ce54b864 100644
--- a/Sources/ShieldSecurity/SecKeyPair.swift
+++ b/Sources/ShieldSecurity/SecKeyPair.swift
@@ -317,10 +317,8 @@ public struct SecKeyPair {
     case bits256 = 32
   }
 
-  /// Encrypt and encode the key pair's private key using PBKDF.
-  ///
-  /// Encrypt the key pair's private key with a password using PBKDF and then encode
-  /// the encrypted key, along with the PBKDF parameters, into an ASN.1 structure.
+  /// Encodes the key pair's private key in PKCS#8 format and then encrypts it using PBKDF and packages
+  /// into PKCS#8 encrypted format.
   ///
   /// With the exported key and original password, ``import(fromData:withPassword:)``
   /// can be used to recover the original `SecKey`.
@@ -328,8 +326,9 @@ public struct SecKeyPair {
   /// - Parameters:
   ///   - password: Password use for key encryption.
   ///   - derivedKeySize: PBKDF target key size.
+  ///   - psuedoRandomAlgorithm: Which psuedo random algorithm should be used with PBKDF.
   ///   - keyDerivationTiming: Time PBKDF function should take to generate encryption key.
-  /// - Returns: Encoded encrypted key and PBKDF paraemters.
+  /// - Returns: Encrypted PKCS#8 encoded private key.
   ///
   public func export(
     password: String,
@@ -392,15 +391,28 @@ public struct SecKeyPair {
     return encryptedPrivateKeyInfoData
   }
 
-  /// Decode and decrypt a previously exported private key.
+  /// Encodes the key pair's private key in PKCS#8 format.
+  ///
+  /// With the exported key and original password, ``import(fromData:withPassword:)``
+  /// can be used to recover the original `SecKey`.
   ///
-  /// Decodes the encrypted key and PBKDF paraameters from the provided ASN.1 data and then decrypts
-  /// the private key. This is the reverse operation of ``export(password:derivedKeyLength:keyDerivationTiming:)``.
+  /// - Returns: Encoded encrypted key and PBKDF paraemters.
+  ///
+  public func export() throws -> Data {
+
+    return try privateKey.encodePKCS8()
+  }
+
+  /// Decrypts an encrypted PKCS#8 encrypted private key and builds a complete key pair.
+  ///
+  /// This is the reverse operation of ``export(password:derivedKeyLength:keyDerivationTiming:)``.
+  ///
+  /// - Note: Only supports PKCS#8's PBES2 sceheme using PBKDF2 for key derivation.
   ///
   /// - Parameters:
   ///   - data: Data for exported private key.
   ///   - password: Password used during key export.
-  /// - Returns: ``SecKeyPair`` for the decoded/decrypted private.
+  /// - Returns: ``SecKeyPair`` for the decrypted & decoded private key.
   ///
   public static func `import`(fromData data: Data, withPassword password: String) throws -> SecKeyPair {
 
@@ -444,9 +456,20 @@ public struct SecKeyPair {
                                                key: importKey,
                                                iv: aesIV)
 
+    return try Self.import(fromData: privateKeyInfoData)
+  }
+
+  /// Decodes a PKCS#8 encoded private key and builds a complete key pair.
+  ///
+  /// - Parameters:
+  ///   - data: Data for exported private key.
+  /// - Returns: ``SecKeyPair`` for the decrypted private key.
+  ///
+  public static func `import`(fromData data: Data) throws -> SecKeyPair {
+
     let privateKeyInfo: PrivateKeyInfo
     do {
-      privateKeyInfo = try ASN1.Decoder.decode(PrivateKeyInfo.self, from: privateKeyInfoData)
+      privateKeyInfo = try ASN1.Decoder.decode(PrivateKeyInfo.self, from: data)
     }
     catch {
       throw SecKeyPair.Error.invalidEncodedPrivateKey
diff --git a/Tests/SecKeyPairTests.swift b/Tests/SecKeyPairTests.swift
index da17ce3f7..b7df44fa7 100644
--- a/Tests/SecKeyPairTests.swift
+++ b/Tests/SecKeyPairTests.swift
@@ -156,7 +156,7 @@ class SecKeyPairTests: XCTestCase {
   }
 #endif
 
-  func testImportExportRSA() throws {
+  func testImportExportEncryptedRSA() throws {
 
     let exportedKeyData = try rsaKeyPair.export(password: "123")
 
@@ -189,7 +189,22 @@ class SecKeyPairTests: XCTestCase {
 
   }
 
-  func testImportExportEC() throws {
+  func testImportExportRSA() throws {
+
+    let exportedKeyData = try rsaKeyPair.export()
+
+    let importedKeyPair = try SecKeyPair.import(fromData: exportedKeyData)
+
+    let plainText = try Random.generate(count: 171)
+
+    let cipherText1 = try rsaKeyPair.publicKey.encrypt(plainText: plainText, padding: .oaep)
+
+    let plainText2 = try importedKeyPair.privateKey.decrypt(cipherText: cipherText1, padding: .oaep)
+
+    XCTAssertEqual(plainText, plainText2)
+  }
+
+  func testImportExportEncryptedEC() throws {
 
     let exportedKeyData = try ecKeyPair.export(password: "123")
 
@@ -198,6 +213,13 @@ class SecKeyPairTests: XCTestCase {
     XCTAssertThrowsError(try SecKeyPair.import(fromData: exportedKeyData, withPassword: "456"))
   }
 
+  func testImportExportEC() throws {
+
+    let exportedKeyData = try ecKeyPair.export()
+
+    _ = try SecKeyPair.import(fromData: exportedKeyData)
+  }
+
   func testCodable() throws {
 
     let rsaData = try JSONEncoder().encode(rsaKeyPair)