//
//  Data+Encryption.swift
//  Function
//
//  Created by mac on 2016/11/9.
//  Copyright © 2016年 mac. All rights reserved.
//

import UIKit

// MARK: RSA encrypt
extension Data {
    fileprivate func rsa_publickey_form_data(keyData: Data) -> SecKey? {
        if let certificate = SecCertificateCreateWithData(kCFAllocatorDefault, keyData as CFData) {
            let policy = SecPolicyCreateBasicX509()
            var trust: SecTrust?
            if SecTrustCreateWithCertificates(certificate, policy, &trust) == errSecSuccess {
                var trustResultType: SecTrustResultType = SecTrustResultType.invalid
                if SecTrustEvaluate(trust!, &trustResultType) == errSecSuccess {
                    return SecTrustCopyPublicKey(trust!)!
                }
            }
        }
        return nil
    }

    fileprivate func rsa_privatekey_from_data(keyData: Data, withPassword password: String) -> SecKey? {
        var privateKey: SecKey? = nil
        let options: [String: String] = [kSecImportExportPassphrase as String: password]
        var items: CFArray?
        if SecPKCS12Import(keyData as CFData, options as CFDictionary, &items) == errSecSuccess {
            if CFArrayGetCount(items) > 0 {
                let d = unsafeBitCast(CFArrayGetValueAtIndex(items, 0), to: CFDictionary.self)
                let k = Unmanaged.passUnretained(kSecImportItemIdentity).toOpaque()
                let v = CFDictionaryGetValue(d, k)
                let secIdentity = unsafeBitCast(v, to: SecIdentity.self)
                if SecIdentityCopyPrivateKey(secIdentity, &privateKey) == errSecSuccess {
                    return privateKey
                }
            }
        }
        return nil
    }
    fileprivate func RSA(operation: String, key: SecKey) -> Data? {
        let key_size = SecKeyGetBlockSize(key)
        var encrypt_bytes = [UInt8](repeating: 0, count: key_size)
        var output_size = key_size
        if operation == "encrypt" {
            if SecKeyEncrypt(key, SecPadding.PKCS1,
                             self.bytes, self.count,
                             &encrypt_bytes, &output_size) == errSecSuccess {
                return Data(bytes: encrypt_bytes, count: output_size)
            }
        } else {
            let stauts = SecKeyDecrypt(key, SecPadding.PKCS1,
                                       self.bytes, self.count,
                                       &encrypt_bytes, &output_size)
            if stauts == errSecSuccess {
                return Data(bytes: UnsafePointer<UInt8>(encrypt_bytes), count: output_size)
            }
        }

        return nil
    }

    func RSAEncryptToData(publicKeyPath: String) -> Data {
        let publicKey = try? Data(contentsOf: URL(fileURLWithPath: publicKeyPath))
        let publickeyData = rsa_publickey_form_data(keyData: publicKey!)
        return RSA(operation: "encrypt", key: publickeyData!)!
    }

    func RSAEncryptToBase64Data(publicKeyPath: String) -> Data {
        return RSAEncryptToData(publicKeyPath: publicKeyPath).base64EncodedData()
    }

    func RSAEncryptToBase64String(publicKeyPath: String) -> String {
        return RSAEncryptToData(publicKeyPath: publicKeyPath).base64EncodedString()
    }

    mutating func RSADecryptFromBase64DataToData(privateKeyPath: String) -> Data {
        self = Data.init(base64Encoded: self)!
        return RSADecryptToData(privateKeyPath: privateKeyPath)
    }
    mutating func RSADecryptFromBase64DataToString(privateKeyPath: String) -> String {
        self = Data.init(base64Encoded: self)!
        return RSADecryptToString(privateKeyPath: privateKeyPath)
    }

    func RSADecryptToData(privateKeyPath: String) -> Data {
        let privateKey = try? Data(contentsOf: URL(fileURLWithPath: privateKeyPath))
        let privateKeyData = rsa_privatekey_from_data(keyData: privateKey!, withPassword: "5995267")
        return RSA(operation: "decrypt", key: privateKeyData!)!
    }
    func RSADecryptToString(privateKeyPath: String) -> String {
         return String(data: RSADecryptToData(privateKeyPath: privateKeyPath), encoding: String.Encoding.utf8)!
    }
}

extension String {

    func RSAEncryptToData(publicKeyPath: String) -> Data {

        return self.myData.RSAEncryptToData(publicKeyPath: publicKeyPath)
    }

    func RSAEncryptToBase64Data(publicKeyPath: String) -> Data {
        return self.myData.RSAEncryptToBase64Data(publicKeyPath: publicKeyPath)
    }

    func RSAEncryptToBase64String(publicKeyPath: String) -> String {
        return self.myData.RSAEncryptToBase64String(publicKeyPath: publicKeyPath)
    }

    func RSADecryptFromBase64StringToData(privateKeyPath: String) -> Data {
        return (Data(base64Encoded: self)?.RSADecryptToData(privateKeyPath: privateKeyPath))!
    }

    func RSADecryptFromBase64StringToString(privateKeyPath: String) -> String {
        return (Data(base64Encoded: self)?.RSADecryptToString(privateKeyPath: privateKeyPath))!
    }

    func RSADecryptToData(privateKeyPath: String) -> Data {
        return self.myData.RSADecryptToData(privateKeyPath: privateKeyPath)
    }
    func RSADecryptToString(privateKeyPath: String) -> String {
        return self.myData.RSADecryptToString(privateKeyPath: privateKeyPath)
    }
}