import Foundation

extension NSData {

    func encryptAES128(key: NSString, iv: NSString) -> NSData? {
        return aes128(method: CCOperation(kCCEncrypt), key: key, iv: iv)
    }

    func decryptAES128(key: NSString, iv: NSString) -> NSData? {
        return aes128(method: CCOperation(kCCDecrypt), key: key, iv: iv)
    }

    private func aes128(method: CCOperation, key: NSString, iv: NSString) -> NSData? {
        let keyPtr = UnsafeMutablePointer<CChar>.allocate(capacity: Int(kCCKeySizeAES128) + 1)

        defer {
            keyPtr.deallocate()
        }

        let ivPtr = iv.utf8String
        bzero(keyPtr, 0)
        key.getCString(keyPtr, maxLength: Int(kCCKeySizeAES128) + 1, encoding: String.Encoding.utf8.rawValue)

        let bufferSize = self.length + kCCBlockSizeAES128
        let buffer = UnsafeMutableRawPointer.allocate(byteCount: bufferSize, alignment: MemoryLayout.alignment(ofValue: CChar.self))

        var numBytesEncrypted = 0

        let cryptStatus = CCCrypt(method, CCAlgorithm(kCCAlgorithmAES128), CCOptions(kCCOptionPKCS7Padding), keyPtr, kCCKeySizeAES128, ivPtr, self.bytes, self.length, buffer, bufferSize, &numBytesEncrypted)

        if cryptStatus == kCCSuccess {
            return NSData(bytesNoCopy: buffer, length: numBytesEncrypted, freeWhenDone: true)
        }

        buffer.deallocate()
        return nil
    }

}