//
/*
 * Copyright (c) 2021 BWI GmbH
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import Foundation
import CommonCrypto


struct CryptoAES {
    private let key: Data
    private let ivSize: Int = kCCBlockSizeAES128
    private let options: CCOptions = CCOptions(kCCOptionPKCS7Padding)
    
    enum Error: Swift.Error {
        case invalidKeySize
        case generateRandomIVFailed
        case encryptionFailed
        case decryptionFailed
    }

    init(key: Data) throws {
        guard key.count == kCCKeySizeAES256 else {
            throw Error.invalidKeySize
        }
        self.key = key
    }
}

private extension CryptoAES {

    func generateRandomIV(for data: inout Data) throws {

        try data.withUnsafeMutableBytes { dataBytes in

            guard let dataBytesBaseAddress = dataBytes.baseAddress else {
                throw Error.generateRandomIVFailed
            }

            let status: Int32 = SecRandomCopyBytes(
                kSecRandomDefault,
                kCCBlockSizeAES128,
                dataBytesBaseAddress
            )

            guard status == 0 else {
                throw Error.generateRandomIVFailed
            }
        }
    }
}

extension CryptoAES: Cryptable {

    func encrypt(_ data: Data) throws -> Data {
        var numberBytesEncrypted = 0
        let bufferSize: Int = ivSize + data.count + kCCBlockSizeAES128
        var buffer = Data(count: bufferSize)
        try generateRandomIV(for: &buffer)

        do {
            try key.withUnsafeBytes { keyBytes in
                try data.withUnsafeBytes { dataToEncryptBytes in
                    try buffer.withUnsafeMutableBytes { bufferBytes in

                        guard let keyBytesBaseAddress = keyBytes.baseAddress,
                            let dataToEncryptBytesBaseAddress = dataToEncryptBytes.baseAddress,
                            let bufferBytesBaseAddress = bufferBytes.baseAddress else {
                                throw Error.encryptionFailed
                        }

                        let cryptStatus: CCCryptorStatus = CCCrypt(
                            CCOperation(kCCEncrypt),
                            CCAlgorithm(kCCAlgorithmAES),
                            options,
                            keyBytesBaseAddress,
                            key.count,
                            bufferBytesBaseAddress,
                            dataToEncryptBytesBaseAddress,
                            dataToEncryptBytes.count,
                            bufferBytesBaseAddress + ivSize,
                            bufferSize,
                            &numberBytesEncrypted
                        )

                        if CCCryptorStatus(kCCSuccess) != cryptStatus {
                            throw Error.encryptionFailed
                        }
                    }
                }
            }

        } catch {
            throw Error.encryptionFailed
        }

        let encryptedData: Data = buffer[..<(numberBytesEncrypted + ivSize)]
        return encryptedData
    }

    func decrypt(_ data: Data) throws -> Data {
        var numberBytesDecrypted: Int = 0
        let bufferSize: Int = data.count - ivSize
        var buffer = Data(count: bufferSize)

        do {
            try key.withUnsafeBytes { keyBytes in
                try data.withUnsafeBytes { dataToDecryptBytes in
                    try buffer.withUnsafeMutableBytes { bufferBytes in

                        guard let keyBytesBaseAddress = keyBytes.baseAddress,
                            let dataToDecryptBytesBaseAddress = dataToDecryptBytes.baseAddress,
                            let bufferBytesBaseAddress = bufferBytes.baseAddress else {
                                throw Error.decryptionFailed
                        }

                        let cryptStatus: CCCryptorStatus = CCCrypt(
                            CCOperation(kCCDecrypt),
                            CCAlgorithm(kCCAlgorithmAES),
                            options,
                            keyBytesBaseAddress,
                            key.count,
                            dataToDecryptBytesBaseAddress,
                            dataToDecryptBytesBaseAddress + ivSize,
                            bufferSize,
                            bufferBytesBaseAddress,
                            bufferSize,
                            &numberBytesDecrypted
                        )

                        if CCCryptorStatus(kCCSuccess) != cryptStatus {
                            throw Error.decryptionFailed
                        }
                    }
                }
            }
        } catch {
            throw Error.encryptionFailed
        }

        return buffer[..<numberBytesDecrypted]
    }
}