//
//  CryptoAES.swift
//  PrivadoVPN
//
//  Created by Zhandos Bolatbekov on 15.01.2021.
//  Copyright © 2021 Privado LLC. All rights reserved.
//

import Foundation
import CommonCrypto

struct CryptoAES {
    private let key: Data
    private let iv: Data
    
    init(key: Data, iv: Data) {
        self.key = key.sha256()
        self.iv = iv
    }
    
    func encrypt(message: String) -> Data? {
        guard let messageData = message.data(using: .utf8) else { return nil }
        return crypt(data: messageData, option: CCOperation(kCCEncrypt))
    }
    
    func decrypt(data: Data) -> String? {
        let decryptedData = crypt(data: data, option: CCOperation(kCCDecrypt))
        return String(bytes: decryptedData, encoding: .utf8)
    }
        
    private func crypt(data: Data, option: CCOperation) -> Data {
        let cryptLength = size_t(data.count + kCCBlockSizeAES128)
        var cryptData = Data(count: cryptLength)

        let keyLength = size_t(kCCKeySizeAES256)
        let options = CCOptions(kCCOptionPKCS7Padding)
        var numBytesEncrypted: size_t = 0
        
        let cryptStatus = cryptData.withUnsafeMutableBytes { cryptBytes in
            data.withUnsafeBytes { dataBytes in
                key.withUnsafeBytes { keyBytes in
                    iv.withUnsafeBytes { ivBytes in
                        CCCrypt(option,
                                CCAlgorithm(kCCAlgorithmAES),
                                options,
                                keyBytes.baseAddress,
                                keyLength,
                                ivBytes.baseAddress,
                                dataBytes.baseAddress,
                                data.count,
                                cryptBytes.baseAddress,
                                cryptLength,
                                &numBytesEncrypted)
                    }
                }
            }
        }

        if UInt32(cryptStatus) == UInt32(kCCSuccess) {
            cryptData.removeSubrange(numBytesEncrypted..<cryptData.count)
        } else {
            print("Error: \(cryptStatus)")
        }

        return cryptData
    }
}

private extension Data {
    func sha256() -> Data {
        var hash = [UInt8](repeating: 0,  count: Int(CC_SHA256_DIGEST_LENGTH))
        self.withUnsafeBytes {
            _ = CC_SHA256($0.baseAddress, CC_LONG(self.count), &hash)
        }
        return Data(hash)
    }
}