RSA Key-Pair generator with Encryption/Decryption support

RSAEncryptor.swift

//Reference: https://stackoverflow.com/questions/53906275/rsa-public-key-created-in-ios-swift-and-exported-as-base64-not-recognized-in-jav

import SwiftyRSA

class RSAKeyEncoding: NSObject {
    
    // ASN.1 identifiers
    private let kASNBitStringIdentifier: UInt8 = 0x03
    private let kASNSequenceIdentifier: UInt8 = 0x30
    
    // ASN.1 AlgorithmIdentfier for RSA encryption containing OID 1 2 840 113549 1 1 1; NULL
    private let kASNAlgorithmIdentifierForRSAEncryption: [UInt8] = [0x30, 0x0d, 0x06, 0x09,
                                                                    0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00]
    
    /**
     * Converts the DER encoding of an RSA public key fetched from the keychain (e.g. by
     * using SecItemCopyMatching) to a format typically used by tools and programming
     * languages outside the iOS ecosystem (e.g. OpenSSL, Java, PHP and Perl). The DER
     * encoding of an RSA public key created by iOS is represented with the ASN.1
     * RSAPublicKey type as defined by PKCS#1. However, many systems outside the Apple
     * ecosystem expect the DER encoding of a key to be represented with the ASN.1
     * SubjectPublicKeyInfo type as defined by X.509. The two types are related in a way
     * that if the SubjectPublicKeyInfo’s algorithm field contains the rsaEncryption object
     * identifier as defined by PKCS#1, the subjectPublicKey field must contain the DER
     * encoding of an RSA public key that is represented with the RSAPublicKey type.
     * Returns the converted DER encoding of provided key.
     */
    func convertToX509EncodedKey(_ rsaPublicKeyData: Data) -> Data {
        
        // Initialize an array that will be filled with bytes of the X509 encoded key
        var derEncodedKeyBytes: [UInt8] = [UInt8](rsaPublicKeyData)
        
        // Insert ASN.1 BIT STRING bytes at the beginning of the array
        derEncodedKeyBytes.insert(0x00, at: 0)
        derEncodedKeyBytes.insert(contentsOf: lengthField(of: derEncodedKeyBytes), at: 0)
        derEncodedKeyBytes.insert(kASNBitStringIdentifier, at: 0)
        
        // Insert ASN.1 AlgorithmIdentifier bytes at the beginning of the array
        derEncodedKeyBytes.insert(contentsOf: kASNAlgorithmIdentifierForRSAEncryption, at: 0)
        
        // Insert ASN.1 SEQUENCE bytes at the beginning of the array
        derEncodedKeyBytes.insert(contentsOf: lengthField(of: derEncodedKeyBytes), at: 0)
        derEncodedKeyBytes.insert(kASNSequenceIdentifier, at: 0)
        
        return Data(bytes: derEncodedKeyBytes)
    }
    
    private func lengthField(of valueField: [UInt8]) -> [UInt8] {
        var length = valueField.count
        
        if length < 128 {
            return [ UInt8(length) ]
        }
        
        // Number of bytes needed to encode the length
        let lengthBytesCount = Int((log2(Double(length)) / 8) + 1)
        
        // First byte in length field encoding the number of remaining bytes in this field
        let firstLengthFieldByte = UInt8(128 + lengthBytesCount)
        
        var lengthField: [UInt8] = []
        for _ in 0..<lengthBytesCount {
            // Take the last 8 bits of length
            let lengthByte = UInt8(length & 0xff)
            // Insert them at the beginning of the array
            lengthField.insert(lengthByte, at: 0)
            // Delete the last 8 bits of length
            length = length >> 8
        }
        
        // Insert firstLengthFieldByte at the beginning of the array
        lengthField.insert(firstLengthFieldByte, at: 0)
        
        return lengthField
    }
}

public class RSAEncryptor: NSObject {
    
    private let publicKeyTag: String
    private let privateKeyTag: String
    private let keySize: Int
    
    var publicKey: SecKey?
    var privateKey: SecKey?
    @objc public let refId = UUID().uuidString
    
    @objc public static let sharedEncryptor: RSAEncryptor = RSAEncryptor()
    
    init(publicKeyTag:String = "com.rsakeypair.publickey", privateKeyTag:String = "com.rsakeypair.privatekey", keySize:Int = 1024) {
        self.publicKeyTag = publicKeyTag
        self.privateKeyTag = privateKeyTag
        self.keySize = keySize
        super.init()
        generateKeypair()
    }
    
    private func generateKeypair() {
        var sanityCheck: OSStatus = noErr
        // Container dictionaries
        var privateKeyAttr = [AnyHashable : Any]()
        var publicKeyAttr = [AnyHashable: Any]()
        var keyPairAttr = [AnyHashable : Any]()
        // Set top level dictionary for the keypair
        keyPairAttr[(kSecAttrKeyType ) as AnyHashable] = (kSecAttrKeyTypeRSA as Any)
        keyPairAttr[(kSecAttrKeySizeInBits as AnyHashable)] = Int(keySize)
        // Set private key dictionary
        privateKeyAttr[(kSecAttrIsPermanent as AnyHashable)] = Int(truncating: true)
        privateKeyAttr[(kSecAttrApplicationTag as AnyHashable)] = privateKeyTag
        // Set public key dictionary.
        publicKeyAttr[(kSecAttrIsPermanent as AnyHashable)] = Int(truncating: true)
        publicKeyAttr[(kSecAttrApplicationTag as AnyHashable)] = publicKeyTag
        
        keyPairAttr[(kSecPrivateKeyAttrs as AnyHashable)] = privateKeyAttr
        keyPairAttr[(kSecPublicKeyAttrs as AnyHashable)] = publicKeyAttr
        sanityCheck = SecKeyGeneratePair((keyPairAttr as CFDictionary), &publicKey, &privateKey)
        if sanityCheck == noErr && publicKey != nil && privateKey != nil {
            print("RSA key pair generation Successfull")
        }
    }
    
    private func appendPrefixSuffixTo(_ string: String, prefix: String, suffix: String) -> String {
        return "\(prefix)\(string)\(suffix)"
    }
    
    @objc public func publicKeyDer() -> String? {
        if let pemKey = self.publicKey?.derFormat() {
            return appendPrefixSuffixTo(pemKey, prefix: "", suffix: "")
        }
        return nil
    }
    
    @objc public func privateKeyPem() -> String? {
        return appendPrefixSuffixTo(self.privateKey?.pemFormat() ?? "NO KEY", prefix: "", suffix: "")
    }
    
    @objc public func encrypt(data: String) -> String? {
        let blockSize = SecKeyGetBlockSize(publicKey!)
        var messageEncrypted = [UInt8](repeating: 0, count: blockSize)
        var messageEncryptedSize = blockSize
        
        var status: OSStatus!
        
        status = SecKeyEncrypt(publicKey!, SecPadding.PKCS1, data, data.count, &messageEncrypted, &messageEncryptedSize)
        
        if status != noErr {
            print("Encryption Error!")
            return nil
        }
        return String(bytes: messageEncrypted, encoding: .utf8)
    }
    
    @objc public func decrypt(data: String) -> String? {
        do {
            let pvtKey = try PrivateKey(reference: privateKey!)
            let encryptedMessage = try EncryptedMessage(base64Encoded: data)
            let clearMsg = try encryptedMessage.decrypted(with:pvtKey, padding: .PKCS1)
            return try clearMsg.string(encoding: .utf8)
        } catch let error {
            print(error.localizedDescription)
        }
        return nil
    }
}

extension SecKey {
    public func derFormat() -> String? {
        var pbError:Unmanaged<CFError>?
        if #available(iOS 10, *) {
            guard let pbData = SecKeyCopyExternalRepresentation(self, &pbError) as Data? else {
                print("error: ", pbError!.takeRetainedValue() as Error)
                return nil
            }
            let publicKeyPemFormat = RSAKeyEncoding().convertToX509EncodedKey(pbData).base64EncodedString()
            print("DER format key: \n",publicKeyPemFormat)
            return publicKeyPemFormat
        } else {
            // On iOS 8/9, we need to add the key again to the keychain with a temporary tag, grab the data,
            // and delete the key again.
            let temporaryTag = UUID().uuidString
            let addParams: [CFString: Any] = [
                kSecValueRef: self as Any,
                kSecReturnData: true,
                kSecClass: kSecClassKey,
                kSecAttrApplicationTag: temporaryTag
            ]
            
            var data: AnyObject?
            let addStatus = SecItemAdd(addParams as CFDictionary, &data)
            guard let unwrappedData = data as? Data else {
                return nil
            }
            
            let deleteParams: [CFString: Any] = [
                kSecClass: kSecClassKey,
                kSecAttrApplicationTag: temporaryTag
            ]
            if addStatus == errSecSuccess {
                _ = SecItemDelete(deleteParams as CFDictionary)
            }
            
            let keyDerFormat = RSAKeyEncoding().convertToX509EncodedKey(unwrappedData).base64EncodedString()
            print("public key: \n", keyDerFormat)
            return keyDerFormat
        }
    }
    
    public func pemFormat() -> String? {
        var pbError:Unmanaged<CFError>?
        if #available(iOS 10, *) {
            guard let pbData = SecKeyCopyExternalRepresentation(self, &pbError) as Data? else {
                print("error: ", pbError!.takeRetainedValue() as Error)
                return nil
            }
            let publicKeyPemFormat = pbData.base64EncodedString()
            print("public key: \n", publicKeyPemFormat)
            return publicKeyPemFormat
        } else {
            // On iOS 8/9, we need to add the key again to the keychain with a temporary tag, grab the data,
            // and delete the key again.
            let temporaryTag = UUID().uuidString
            let addParams: [CFString: Any] = [
                kSecValueRef: self as Any,
                kSecReturnData: true,
                kSecClass: kSecClassKey,
                kSecAttrApplicationTag: temporaryTag
            ]
            
            var data: AnyObject?
            let addStatus = SecItemAdd(addParams as CFDictionary, &data)
            guard let unwrappedData = data as? Data else {
                return nil
            }
            
            let deleteParams: [CFString: Any] = [
                kSecClass: kSecClassKey,
                kSecAttrApplicationTag: temporaryTag
            ]
            if addStatus == errSecSuccess {
                _ = SecItemDelete(deleteParams as CFDictionary)
            }
            
            let publicKeyPemFormat = unwrappedData.base64EncodedString()
            print("public key: \n", publicKeyPemFormat)
            return publicKeyPemFormat
        }
    }
    
    public func asBase64() -> String? {
        var dataPtr: CFTypeRef?
        let query: [String:Any] = [
            kSecClass as String: kSecClassKey,
            kSecAttrApplicationTag as String: "com.wibmo.cardsign.publickey", // Same unique tag here
            kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
            kSecReturnData as String: kCFBooleanTrue
        ]
        let result = SecItemCopyMatching(query as CFDictionary, &dataPtr)
        
        switch (result, dataPtr) {
        case (errSecSuccess, .some(let data)):
            
            // convert to X509 encoded key
            let convertedData = RSAKeyEncoding().convertToX509EncodedKey(data as! Data)
            
            // convert to Base64 string
            let base64PublicKey = convertedData.base64EncodedString(options: [])
            return base64PublicKey
        default:
            return nil
//            throw CryptoError.keyConversionError
        }
    }
}