EasyAES

Cryptor.swift

import Foundation
import CommonCrypto

internal enum CryptoError: Error {
  case emptyData
  case invalidData
  case invalidKey
  case fail(status: CCStatus)
}

internal func encryptor(withKey key: Data) -> (Data) throws -> Data {
  return { (input: Data) throws -> Data in
    guard !input.isEmpty else {
      throw CryptoError.emptyData
    }
    guard key.count > 4 else {
      throw CryptoError.invalidKey
    }
    let iv: Data = randomIv()
    return try iv + crypt(input, key: key, iv: iv, operation: CCOperation(kCCEncrypt))
  }
}

internal func decryptor(withKey key: Data) -> (Data) throws -> Data {
  return { (input: Data) throws -> Data in
    guard !input.isEmpty else {
      throw CryptoError.emptyData
    }
    guard key.count > 4 else {
      throw CryptoError.invalidKey
    }
    guard input.count > kCCBlockSizeAES128 else {
      throw CryptoError.invalidData
    }
    let iv: Data = input.prefix(upTo: kCCBlockSizeAES128)
    let input: Data = input.suffix(from: kCCBlockSizeAES128)
    return try crypt(input, key: key, iv: iv, operation: CCOperation(kCCDecrypt))
  }
}

internal func randomKey() -> Data {
  return randomData(with: kCCKeySizeAES256)
}

private func randomIv() -> Data {
  return randomData(with: kCCBlockSizeAES128)
}

private func randomData(with length: Int) -> Data {
  precondition(length > 0)
  var data = Data(count: length)
  guard (data.withUnsafeMutableBytes { (dataBytes: UnsafeMutableRawBufferPointer) in
    return SecRandomCopyBytes(kSecRandomDefault, length, dataBytes.baseAddress!) == errSecSuccess
  }) else {
    data.withUnsafeMutableBytes { arc4random_buf($0.baseAddress, length) }
    return data
  }
  return data
}

private func processed(key: Data, iv: Data) throws -> Data {
  let salt: Data
  if iv.first?.isMultiple(of: 2) ?? false {
    salt = iv[iv.startIndex ..< iv.startIndex.advanced(by: iv.count / 2)]
  } else {
    salt = iv[iv.startIndex.advanced(by: iv.count / 2) ..< iv.endIndex]
  }
  var dataOut: Data = .init(count: Int(CC_SHA256_DIGEST_LENGTH))
  let status = dataOut.withUnsafeMutableBytes { (outputBytes: UnsafeMutableRawBufferPointer) -> CCCryptorStatus in
    salt.withUnsafeBytes { (saltBytes: UnsafeRawBufferPointer) -> CCCryptorStatus in
      key.withUnsafeBytes { (keyBytes: UnsafeRawBufferPointer) -> CCCryptorStatus in
        CCKeyDerivationPBKDF(
          CCPBKDFAlgorithm(kCCPBKDF2),
          keyBytes.bindMemory(to: Int8.self).baseAddress,
          keyBytes.count,
          saltBytes.bindMemory(to: UInt8.self).baseAddress,
          saltBytes.count,
          CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA256),
          10_000,
          outputBytes.bindMemory(to: UInt8.self).baseAddress,
          outputBytes.count
        )
      }
    }
  }
  guard status == kCCSuccess else {
    throw CryptoError.fail(status: status)
  }
  return dataOut
}

private func crypt(_ input: Data, key: Data, iv: Data, operation: CCOperation) throws -> Data {
  precondition(!key.isEmpty)
  precondition(!iv.isEmpty)
  let key = try processed(key: key, iv: iv)
  guard key.count == kCCKeySizeAES256 else {
    throw CryptoError.invalidKey
  }
  var dataOutCount = 0
  var output: Data = .init(count: input.count + kCCBlockSizeAES128)
  let status = output.withUnsafeMutableBytes { (outputBytes: UnsafeMutableRawBufferPointer) -> CCCryptorStatus in
    input.withUnsafeBytes { (inputBytes: UnsafeRawBufferPointer) -> CCCryptorStatus in
      iv.withUnsafeBytes { (ivBytes: UnsafeRawBufferPointer) -> CCCryptorStatus in
        key.withUnsafeBytes { (keyBytes: UnsafeRawBufferPointer) -> CCCryptorStatus in
          CCCrypt(
            operation,
            CCPBKDFAlgorithm(kCCAlgorithmAES),
            CCOptions(kCCOptionPKCS7Padding),
            keyBytes.baseAddress,
            kCCKeySizeAES256,
            ivBytes.baseAddress,
            inputBytes.baseAddress,
            inputBytes.count,
            outputBytes.baseAddress,
            outputBytes.count,
            &dataOutCount
          )
        }
      }
    }
  }
  guard status == kCCSuccess else {
    throw CryptoError.fail(status: status)
  }
  output = output.prefix(upTo: dataOutCount)
  precondition(output != input)
  return output
}

Example:

let key: Data = randomKey()
let someData: Data = .init(repeating: 0b0101, count: 64)
let encrypted = try encryptor(withKey:key)(someData)
let decrypted = try decryptor(withKey:key)(encrypted)