From e62f4714e806a95acd75c9ebd3c607613c18b99e Mon Sep 17 00:00:00 2001
From: Mingshen Sun <bob@mssun.me>
Date: Mon, 13 Apr 2020 10:25:01 -0700
Subject: [PATCH] Encryption support for multi keys

---
 passKit/Models/PasswordStore.swift          | 4 +++-
 passKitTests/Crypto/PGPAgentTest.swift      | 6 +++---
 passKitTests/Models/PasswordStoreTest.swift | 5 +++++
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/passKit/Models/PasswordStore.swift b/passKit/Models/PasswordStore.swift
index ab3c749..c6e5923 100644
--- a/passKit/Models/PasswordStore.swift
+++ b/passKit/Models/PasswordStore.swift
@@ -711,7 +711,9 @@ public class PasswordStore {
     }
     
     public func encrypt(password: Password) throws -> Data {
-        return try PGPAgent.shared.encrypt(plainData: password.plainData, keyID: "")
+        let encryptedDataPath = storeURL.appendingPathComponent(password.url.path)
+        let keyID = findGPGID(from: encryptedDataPath)
+        return try PGPAgent.shared.encrypt(plainData: password.plainData, keyID: keyID)
     }
     
     public func removeGitSSHKeys() {
diff --git a/passKitTests/Crypto/PGPAgentTest.swift b/passKitTests/Crypto/PGPAgentTest.swift
index 5baf3ac..0efd17b 100644
--- a/passKitTests/Crypto/PGPAgentTest.swift
+++ b/passKitTests/Crypto/PGPAgentTest.swift
@@ -31,11 +31,11 @@ class PGPAgentTest: XCTestCase {
         super.tearDown()
     }
 
-    func basicEncryptDecrypt(using pgpAgent: PGPAgent, keyID: String, requestPassphrase: () -> String = requestPGPKeyPassphrase, encryptInArmored: Bool = true, encryptInArmoredNow: Bool = true) throws -> Data? {
+    func basicEncryptDecrypt(using pgpAgent: PGPAgent, keyID: String, encryptKeyID: String? = nil, requestPassphrase: () -> String = requestPGPKeyPassphrase, encryptInArmored: Bool = true, encryptInArmoredNow: Bool = true) throws -> Data? {
         passKit.Defaults.encryptInArmored = encryptInArmored
         let encryptedData = try pgpAgent.encrypt(plainData: testData, keyID: keyID)
         passKit.Defaults.encryptInArmored = encryptInArmoredNow
-        return try pgpAgent.decrypt(encryptedData: encryptedData, keyID: keyID, requestPGPKeyPassphrase: requestPassphrase)
+        return try pgpAgent.decrypt(encryptedData: encryptedData, keyID: encryptKeyID ?? keyID, requestPGPKeyPassphrase: requestPassphrase)
     }
 
     func testMultiKeys() throws {
@@ -104,7 +104,7 @@ class PGPAgentTest: XCTestCase {
     func testIncompatibleKeyTypes() throws {
         try importKeys(ED25519.publicKey, RSA2048.privateKey)
         XCTAssert(pgpAgent.isPrepared)
-        XCTAssertThrowsError(try basicEncryptDecrypt(using: pgpAgent, keyID: RSA2048.fingerprint)) {
+        XCTAssertThrowsError(try basicEncryptDecrypt(using: pgpAgent, keyID: ED25519.fingerprint, encryptKeyID: RSA2048.fingerprint)) {
             XCTAssertEqual($0 as! AppError, AppError.KeyExpiredOrIncompatible)
         }
     }
diff --git a/passKitTests/Models/PasswordStoreTest.swift b/passKitTests/Models/PasswordStoreTest.swift
index 1af9773..93377df 100644
--- a/passKitTests/Models/PasswordStoreTest.swift
+++ b/passKitTests/Models/PasswordStoreTest.swift
@@ -55,6 +55,11 @@ class PasswordStoreTest: XCTestCase {
         let work = try decrypt(passwordStore: passwordStore, path: "work/github.com.gpg", passphrase: "passforios")
         XCTAssertEqual(work.plainText, "passwordforwork\n")
 
+        let testPassword = Password(name: "test", url: URL(string: "test.gpg")!, plainText: "testpassword")
+        let testPasswordEntity = try passwordStore.add(password: testPassword)!
+        let testPasswordPlain = try passwordStore.decrypt(passwordEntity: testPasswordEntity, requestPGPKeyPassphrase: { "passforios" } )!
+        XCTAssertEqual(testPasswordPlain.plainText, "testpassword")
+
         passwordStore.erase()
     }