passforios/passKit/Crypto/GopenPGPInterface.swift

//
//  GopenPGPInterface.swift
//  passKit
//
//  Created by Danny Moesch on 08.09.19.
//  Copyright © 2019 Bob Sun. All rights reserved.
//

import Gopenpgp

struct GopenPGPInterface: PGPInterface {
    private static let errorMapping: [String: Error] = [
        "gopenpgp: error in unlocking key: openpgp: invalid data: private key checksum failure": AppError.wrongPassphrase,
        "gopenpgp: error in reading message: openpgp: incorrect key": AppError.keyExpiredOrIncompatible,
    ]

    private var publicKeys: [String: CryptoKey] = [:]
    private var privateKeys: [String: CryptoKey] = [:]
    private var privateSubkeyToKeyIDMapping: [String: String] = [:] // value is the key in privateKeys map

    init(publicArmoredKey: String, privateArmoredKey: String) throws {
        let pubKeys = extractKeysFromArmored(str: publicArmoredKey)
        let prvKeys = extractKeysFromArmored(str: privateArmoredKey)

        for key in pubKeys {
            var error: NSError?
            guard let cryptoKey = CryptoNewKeyFromArmored(key, &error) else {
                guard error == nil else {
                    throw error!
                }
                throw AppError.keyImport
            }
            publicKeys[cryptoKey.getFingerprint().lowercased()] = cryptoKey
        }

        for key in prvKeys {
            var error: NSError?
            guard let cryptoKey = CryptoNewKeyFromArmored(key, &error) else {
                guard error == nil else {
                    throw error!
                }
                throw AppError.keyImport
            }

            let keyID = cryptoKey.getFingerprint().lowercased()
            privateKeys[keyID] = cryptoKey

            guard let subkeyIDsJSON = HelperPassGetHexSubkeyIDsJSON(cryptoKey) else {
                guard error == nil else {
                    throw error!
                }
                throw AppError.keyImport
            }
            do {
                let subkeyIDs = try JSONDecoder().decode([String].self, from: subkeyIDsJSON)
                for subkeyID in subkeyIDs {
                    privateSubkeyToKeyIDMapping[subkeyID] = keyID
                }
            } catch {
                throw AppError.keyImport
            }
        }
    }

    func extractKeysFromArmored(str: String) -> [String] {
        var keys: [String] = []
        var key = ""
        for line in str.splitByNewline() {
            if line.trimmed.uppercased().hasPrefix("-----BEGIN PGP") {
                key = ""
                key += line
            } else if line.trimmed.uppercased().hasPrefix("-----END PGP") {
                key += line
                keys.append(key)
            } else {
                key += line
            }
            key += "\n"
        }
        return keys
    }

    func containsPublicKey(with keyID: String) -> Bool {
        publicKeys.keys.contains { key in key.hasSuffix(keyID.lowercased()) }
    }

    func containsPrivateKey(with keyID: String) -> Bool {
        privateKeys.keys.contains { key in key.hasSuffix(keyID.lowercased()) }
    }

    func decrypt(encryptedData: Data, keyIDHint: String?, passPhraseForKey: @escaping (String) -> String) throws -> Data? {
        let message = createPGPMessage(from: encryptedData)
        guard let message else {
            throw AppError.decryption
        }

        guard let privateKey: CryptoKey = try findDecryptionKey(message: message, keyIDHint: keyIDHint) else {
            throw AppError.decryption
        }

        do {
            var isLocked: ObjCBool = false
            try privateKey.isLocked(&isLocked)
            var unlockedKey: CryptoKey!
            if isLocked.boolValue {
                let passphrase = passPhraseForKey(privateKey.getFingerprint())
                unlockedKey = try privateKey.unlock(passphrase.data(using: .utf8))
            } else {
                unlockedKey = privateKey
            }
            var error: NSError?

            guard let keyRing = CryptoNewKeyRing(unlockedKey, &error) else {
                guard error == nil else {
                    throw error!
                }
                throw AppError.decryption
            }

            return try keyRing.decrypt(message, verifyKey: nil, verifyTime: 0).data
        } catch {
            throw Self.errorMapping[error.localizedDescription, default: error]
        }
    }

    func encryptWithAllKeys(plainData: Data) throws -> Data {
        let keyIDs = publicKeys.keys.filter { key in privateKeys.keys.contains(key) }
        return try encrypt(plainData: plainData, keyIDs: keyIDs)
    }

    func encrypt(plainData: Data, keyIDs: [String]) throws -> Data {
        let keys: [CryptoKey] = try keyIDs.map { keyID in
            guard let key = publicKeys.first(where: { key, _ in key.hasSuffix(keyID.lowercased()) })?.value else {
                throw AppError.pgpPublicKeyNotFound(keyID: keyID)
            }
            return key
        }
        guard let firstKey = keys.first else {
            throw AppError.encryption
        }
        let otherKeys = keys.dropFirst()

        var error: NSError?
        guard let keyRing = CryptoNewKeyRing(firstKey, &error) else {
            guard error == nil else {
                throw error!
            }
            throw AppError.encryption
        }
        do {
            try otherKeys.forEach { key in
                try keyRing.add(key)
            }
        } catch {
            throw AppError.encryption
        }

        let encryptedData = try keyRing.encrypt(CryptoNewPlainMessage(plainData.mutable as Data), privateKey: nil)
        if Defaults.encryptInArmored {
            var error: NSError?
            let armor = encryptedData.getArmored(&error)
            guard error == nil else {
                throw error!
            }
            return armor.data(using: .ascii)!
        }
        return encryptedData.getBinary()!
    }

    func getKeyIDs(type: PGPKey) -> [String] {
        switch type {
        case .PUBLIC:
            return publicKeys.keys.map { $0.uppercased() }
        case .PRIVATE:
            return privateKeys.keys.map { $0.uppercased() }
        }
    }

    func getShortKeyIDs(type: PGPKey) -> [String] {
        getKeyIDs(type: type).map { $0.suffix(8).uppercased() }
    }

    private func findDecryptionKey(message: CryptoPGPMessage, keyIDHint: String?) throws -> CryptoKey? {
        var keyIDCandidates: any Collection<String> = privateKeys.keys
        do {
            if let encryptionKeysJSON = message.getHexEncryptionKeyIDsJson() {
                // these are the subkeys (encryption keys), not the primaries keys (whose fingerprints we have in the privateKeys map),
                // so we need to map them back to the primary keyIDs using privateSubkeyToKeyIDMapping
                let validSubkeys = try JSONDecoder().decode([String].self, from: encryptionKeysJSON)
                let validKeyIDs = validSubkeys.compactMap { privateSubkeyToKeyIDMapping[$0] }
                if #available(iOSApplicationExtension 16.0, *) {
                    assert(validKeyIDs.isEmpty || !Set(keyIDCandidates).isDisjoint(with: validKeyIDs))
                }
                keyIDCandidates = validKeyIDs
            }
        } catch {
            // fall back to legacy approach of trying first in privateKeys (or preferring hint)
        }

        if let keyIDHint {
            keyIDCandidates = keyIDCandidates.filter { key in key.hasSuffix(keyIDHint.lowercased()) }
        }
        guard let selectedKeyID = keyIDCandidates.first else {
            throw keyIDHint != nil ? AppError.keyExpiredOrIncompatible : AppError.decryption
        }

        return privateKeys[selectedKeyID]
    }
}

public func createPGPMessage(from encryptedData: Data) -> CryptoPGPMessage? {
    // Important note:
    // Even if Defaults.encryptInArmored is true now, it could be different during the encryption.
    var error: NSError?
    let message = CryptoNewPGPMessageFromArmored(String(data: encryptedData, encoding: .ascii), &error)
    if error == nil {
        return message
    }
    return CryptoNewPGPMessage(encryptedData.mutable as Data)
}