office_encrypt

defmodule OfficeEncrypt do
  @moduledoc """
  Encrypts Office Open XML documents (.xlsx, .docx, .pptx) with a password.

  Implements ECMA-376 Agile Encryption using AES-256-CBC with SHA-512 key
  derivation, producing an OLE2 compound document that Microsoft Office and
  LibreOffice can open with the correct password.
  """

  alias CompoundFile.Writer

  @static_file_data %{
                      "\x06DataSpaces/DataSpaceMap" =>
                        "CAAAAAEAAABoAAAAAQAAAAAAAAAgAAAARQBuAGMAcgB5AHAAdABlAGQAUABhAGMAawBhAGcAZQAyAAAAUwB0AHIAbwBuAGcARQBuAGMAcgB5AHAAdABpAG8AbgBEAGEAdABhAFMAcABhAGMAZQAAAA==",
                      "\x06DataSpaces/Version" =>
                        "PAAAAE0AaQBjAHIAbwBzAG8AZgB0AC4AQwBvAG4AdABhAGkAbgBlAHIALgBEAGEAdABhAFMAcABhAGMAZQBzAAEAAAABAAAAAQAAAA==",
                      "\x06DataSpaces/DataSpaceInfo/StrongEncryptionDataSpace" =>
                        "CAAAAAEAAAAyAAAAUwB0AHIAbwBuAGcARQBuAGMAcgB5AHAAdABpAG8AbgBUAHIAYQBuAHMAZgBvAHIAbQAAAA==",
                      "\x06DataSpaces/TransformInfo/StrongEncryptionTransform/\x06Primary" =>
                        "WAAAAAEAAABMAAAAewBGAEYAOQBBADMARgAwADMALQA1ADYARQBGAC0ANAA2ADEAMwAtAEIARABEADUALQA1AEEANAAxAEMAMQBEADAANwAyADQANgB9AE4AAABNAGkAYwByAG8AcwBvAGYAdAAuAEMAbwBuAHQAYQBpAG4AZQByAC4ARQBuAGMAcgB5AHAAdABpAG8AbgBUAHIAYQBuAHMAZgBvAHIAbQAAAAEAAAABAAAAAQAAAAAAAAAAAAAAAAAAAAQAAAA="
                    }
                    |> Map.new(fn {k, v} -> {k, Base.decode64!(v)} end)

  @salt_size 16
  @block_size 16
  @hash_size 64
  @key_bits 256

  @verifier_hash_input_block_key <<0xFE, 0xA7, 0xD2, 0x76, 0x3B, 0x4B, 0x9E, 0x79>>
  @encrypted_verifier_hashvalue_block_key <<0xD7, 0xAA, 0x0F, 0x6D, 0x30, 0x61, 0x34, 0x4E>>
  @encrypted_key_value_block_key <<0x14, 0x6E, 0x0B, 0xE7, 0xAB, 0xAC, 0xD0, 0xD6>>
  @hmac_key_block <<0x5F, 0xB2, 0xAD, 0x01, 0x0C, 0xB9, 0xE1, 0xF6>>
  @hmac_value_block <<0xA0, 0x67, 0x7F, 0x02, 0xB2, 0x2C, 0x84, 0x33>>

  @doc """
  Encrypts an Office Open XML document with the given password.

  The `data` argument should be the raw bytes of an `.xlsx`, `.docx`, or `.pptx`
  file. Returns `{:ok, binary}` on success.

  ## Options

    * `:spin_count` - number of hash iterations for key derivation (default: 100,000).
      Higher values are slower but more resistant to brute-force attacks.

  ## Example

      xlsx_data = File.read!("spreadsheet.xlsx")
      {:ok, encrypted} = OfficeEncrypt.encrypt(xlsx_data, "secret")
      File.write!("spreadsheet_encrypted.xlsx", encrypted)
  """
  @spec encrypt(data :: binary, password :: String.t(), opts :: [spin_count: pos_integer]) ::
          {:ok, binary} | {:error, :empty | :file_size_limit_exceeded | :filename_too_long}
  def encrypt(data, password, opts \\ []) when is_binary(data) and is_binary(password) do
    spin_count = Keyword.get(opts, :spin_count, 100_000)

    {encryption_info, encrypted_package} = perform_encryption(data, password, spin_count)
    create_compound_document(encryption_info, encrypted_package)
  end

  defp perform_encryption(xlsx_data, password, spin_count) do
    key_salt = :crypto.strong_rand_bytes(@salt_size)

    pwd_hash = derive_password_hash(password, key_salt, spin_count)

    verifier_key = derive_block_key(pwd_hash, @verifier_hash_input_block_key)
    verifier = :crypto.strong_rand_bytes(@salt_size)
    encrypted_verifier = aes_encrypt(verifier, verifier_key, key_salt)

    hash_key = derive_block_key(pwd_hash, @encrypted_verifier_hashvalue_block_key)
    verifier_hash = :crypto.hash(:sha512, verifier)
    encrypted_hash = aes_encrypt(verifier_hash, hash_key, key_salt)

    data_key = :crypto.strong_rand_bytes(@salt_size) |> pad_to_size(div(@key_bits, 8), 0x36)
    key_value_key = derive_block_key(pwd_hash, @encrypted_key_value_block_key)
    encrypted_key_value = aes_encrypt(data_key, key_value_key, key_salt)

    package_salt = :crypto.strong_rand_bytes(@salt_size)
    encrypted_package = encrypt_package(xlsx_data, data_key, package_salt)

    encryption_info =
      build_encryption_info(
        key_salt,
        package_salt,
        spin_count,
        encrypted_verifier,
        encrypted_hash,
        encrypted_key_value,
        data_key,
        encrypted_package
      )

    {encryption_info, encrypted_package}
  end

  defp create_compound_document(encryption_info, encrypted_package) do
    document =
      Writer.new()
      |> Writer.add_file("EncryptionInfo", encryption_info)
      |> Writer.add_file("EncryptedPackage", encrypted_package)

    Enum.reduce(@static_file_data, document, fn {name, data}, document ->
      Writer.add_file(document, name, data)
    end)
    |> Writer.render()
  end

  defp derive_password_hash(password, salt, iterations) do
    utf16_password = :unicode.characters_to_binary(password, :utf8, {:utf16, :little})
    initial_hash = :crypto.hash(:sha512, salt <> utf16_password)

    Enum.reduce(0..(iterations - 1), initial_hash, fn i, hash ->
      :crypto.hash(:sha512, <<i::little-32>> <> hash)
    end)
  end

  defp derive_block_key(base_hash, block_key) do
    :crypto.hash(:sha512, base_hash <> block_key) |> binary_part(0, div(@key_bits, 8))
  end

  defp derive_iv(salt, block_key) when is_integer(block_key) do
    derive_iv(salt, <<block_key::little-32>>)
  end

  defp derive_iv(salt, block_key) when is_binary(block_key) do
    :crypto.hash(:sha512, salt <> block_key) |> binary_part(0, @salt_size)
  end

  defp aes_encrypt(plaintext, key, iv) do
    :crypto.crypto_one_time(:aes_256_cbc, key, iv, plaintext, true)
  end

  # This appears to be different from PKCS#7 in that it does not pad if the data is already a
  # multiple of the block size
  defp pad_to_block_boundary(data, block_size \\ @block_size) do
    case rem(byte_size(data), block_size) do
      0 ->
        data

      rem ->
        padding_length = block_size - rem
        data <> :binary.copy(<<padding_length>>, padding_length)
    end
  end

  defp encrypt_package(xlsx_data, key, salt) do
    encrypted_data =
      xlsx_data
      |> chunk_data()
      |> Enum.with_index()
      |> Enum.map(fn {block, index} ->
        iv = derive_iv(salt, index)
        block |> pad_to_block_boundary() |> aes_encrypt(key, iv)
      end)
      |> IO.iodata_to_binary()

    <<byte_size(xlsx_data)::little-64, encrypted_data::binary>>
  end

  defp chunk_data(data, acc \\ [])
  defp chunk_data(<<>>, acc), do: Enum.reverse(acc)

  defp chunk_data(data, acc) when byte_size(data) <= 4096 do
    Enum.reverse([data | acc])
  end

  defp chunk_data(<<chunk::binary-size(4096), rest::binary>>, acc) do
    chunk_data(rest, [chunk | acc])
  end

  defp build_encryption_info(
         key_salt,
         package_salt,
         spin_count,
         enc_verifier,
         enc_hash,
         enc_key,
         data_key,
         enc_package
       ) do
    hmac_salt = :crypto.strong_rand_bytes(@hash_size)
    hmac_value = :crypto.mac(:hmac, :sha512, hmac_salt, enc_package)

    enc_hmac_key = aes_encrypt(hmac_salt, data_key, derive_iv(package_salt, @hmac_key_block))
    enc_hmac_value = aes_encrypt(hmac_value, data_key, derive_iv(package_salt, @hmac_value_block))

    xml_content =
      """
      <?xml version="1.0" encoding="UTF-8" standalone="yes"?>
      <encryption xmlns="http://schemas.microsoft.com/office/2006/encryption" xmlns:p="http://schemas.microsoft.com/office/2006/keyEncryptor/password">
        <keyData saltSize="#{@salt_size}" blockSize="#{@block_size}" keyBits="#{@key_bits}" hashSize="#{@hash_size}" cipherAlgorithm="AES" cipherChaining="ChainingModeCBC" hashAlgorithm="SHA512" saltValue="#{Base.encode64(package_salt)}"/>
        <dataIntegrity encryptedHmacKey="#{Base.encode64(enc_hmac_key)}" encryptedHmacValue="#{Base.encode64(enc_hmac_value)}"/>
        <keyEncryptors>
          <keyEncryptor uri="http://schemas.microsoft.com/office/2006/keyEncryptor/password">
            <p:encryptedKey spinCount="#{spin_count}" saltSize="#{@salt_size}" blockSize="#{@block_size}" keyBits="#{@key_bits}" hashSize="#{@hash_size}" cipherAlgorithm="AES" cipherChaining="ChainingModeCBC" hashAlgorithm="SHA512" saltValue="#{Base.encode64(key_salt)}" encryptedVerifierHashInput="#{Base.encode64(enc_verifier)}" encryptedVerifierHashValue="#{Base.encode64(enc_hash)}" encryptedKeyValue="#{Base.encode64(enc_key)}"/>
          </keyEncryptor>
        </keyEncryptors>
      </encryption>
      """
      |> String.trim()

    <<4::little-16, 4::little-16, 0x40::little-32, xml_content::binary>>
  end

  defp pad_to_size(data, target_size, _byte) when byte_size(data) >= target_size, do: data

  defp pad_to_size(data, target_size, byte) do
    padding_needed = target_size - byte_size(data)
    data <> :binary.copy(<<byte>>, padding_needed)
  end
end