passforios/passKit/Crypto/PGPAgent.swift

//
//  PGPAgent.swift
//  passKit
//
//  Created by Yishi Lin on 2019/7/17.
//  Copyright © 2019 Bob Sun. All rights reserved.
//

public class PGPAgent {
    public static let shared = PGPAgent()

    private let keyStore: KeyStore
    private var pgpInterface: PGPInterface?
    private var latestDecryptStatus = true

    public init(keyStore: KeyStore = AppKeychain.shared) {
        self.keyStore = keyStore
    }

    init(keyStore: KeyStore, pgpInterface: PGPInterface) {
        self.keyStore = keyStore
        self.pgpInterface = pgpInterface
    }

    public func initKeys() throws {
        guard let publicKey: String = keyStore.get(for: PGPKey.PUBLIC.getKeychainKey()),
              let privateKey: String = keyStore.get(for: PGPKey.PRIVATE.getKeychainKey()) else {
            pgpInterface = nil
            throw AppError.keyImport
        }
        do {
            pgpInterface = try GopenPGPInterface(publicArmoredKey: publicKey, privateArmoredKey: privateKey)
        } catch {
            pgpInterface = try ObjectivePGPInterface(publicArmoredKey: publicKey, privateArmoredKey: privateKey)
        }
    }

    public func uninitKeys() {
        pgpInterface = nil
    }

    public func isInitialized() -> Bool {
        pgpInterface != nil
    }

    public func getKeyID() throws -> [String] {
        try checkAndInit()
        return pgpInterface?.keyID ?? []
    }

    public func getShortKeyID() throws -> [String] {
        try checkAndInit()
        return pgpInterface?.shortKeyID.sorted() ?? []
    }

    public func decrypt(encryptedData: Data, keyID: String, requestPGPKeyPassphrase: @escaping (String) -> String) throws -> Data? {
        // Init keys.
        try checkAndInit()
        guard let pgpInterface else {
            throw AppError.decryption
        }

        var keyID = keyID
        if !pgpInterface.containsPrivateKey(with: keyID) {
            if pgpInterface.keyID.count == 1 {
                keyID = pgpInterface.keyID.first!
            } else {
                throw AppError.pgpPrivateKeyNotFound(keyID: keyID)
            }
        }

        // Remember the previous status and set the current status
        let previousDecryptStatus = latestDecryptStatus
        latestDecryptStatus = false

        // Get the PGP key passphrase.
        let providePassPhraseForKey = { (selectedKeyID: String) -> String in
            if previousDecryptStatus == false {
                return requestPGPKeyPassphrase(selectedKeyID)
            }
            return self.keyStore.get(for: AppKeychain.getPGPKeyPassphraseKey(keyID: selectedKeyID)) ?? requestPGPKeyPassphrase(selectedKeyID)
        }
        // Decrypt.
        guard let result = try pgpInterface.decrypt(encryptedData: encryptedData, keyIDHint: keyID, passPhraseForKey: providePassPhraseForKey) else {
            return nil
        }
        // The decryption step has succeed.
        latestDecryptStatus = true
        return result
    }

    public func encrypt(plainData: Data, keyID: String) throws -> Data {
        try checkAndInit()
        guard let pgpInterface else {
            throw AppError.encryption
        }
        var keyID = keyID
        if !pgpInterface.containsPublicKey(with: keyID) {
            if pgpInterface.keyID.count == 1 {
                keyID = pgpInterface.keyID.first!
            } else {
                throw AppError.pgpPublicKeyNotFound(keyID: keyID)
            }
        }
        return try pgpInterface.encrypt(plainData: plainData, keyID: keyID)
    }

    public func decrypt(encryptedData: Data, requestPGPKeyPassphrase: @escaping (String) -> String) throws -> Data? {
        // Remember the previous status and set the current status
        let previousDecryptStatus = latestDecryptStatus
        latestDecryptStatus = false
        // Init keys.
        try checkAndInit()
        // Get the PGP key passphrase.
        let providePassPhraseForKey = { (selectedKeyID: String) -> String in
            if previousDecryptStatus == false {
                return requestPGPKeyPassphrase(selectedKeyID)
            }
            return self.keyStore.get(for: AppKeychain.getPGPKeyPassphraseKey(keyID: selectedKeyID)) ?? requestPGPKeyPassphrase(selectedKeyID)
        }
        // Decrypt.
        guard let result = try pgpInterface!.decrypt(encryptedData: encryptedData, keyIDHint: nil, passPhraseForKey: providePassPhraseForKey) else {
            return nil
        }
        // The decryption step has succeed.
        latestDecryptStatus = true
        return result
    }

    public func encrypt(plainData: Data) throws -> Data {
        try checkAndInit()
        guard let pgpInterface else {
            throw AppError.encryption
        }
        return try pgpInterface.encrypt(plainData: plainData, keyID: nil)
    }

    public var isPrepared: Bool {
        keyStore.contains(key: PGPKey.PUBLIC.getKeychainKey())
            && keyStore.contains(key: PGPKey.PRIVATE.getKeychainKey())
    }

    private func checkAndInit() throws {
        if pgpInterface == nil || !keyStore.contains(key: Globals.pgpKeyPassphrase) {
            try initKeys()
        }
    }
}