defmodule RNS.Crypto.Fernet do
  @moduledoc """
  This module provides a slightly modified implementation of the Fernet spec found at:
  https://github.com/fernet/spec/blob/master/Spec.md

  According to the spec, a Fernet token includes a one byte VERSION and eight byte TIMESTAMP field
  at the start of each token. These fields are not relevant to Reticulum. They are therefore
  stripped from this implementation, since they incur overhead and leak initiator metadata.

  Lazily copied from https://github.com/Sgiath/reticulum
  """
  @type t() :: tuple()
  @type key() :: <<_::256>>
  @type token() :: <<_::512, _::_*8>>

  alias RNS.Crypto

  require Record
  Record.defrecord(:fernet, encryption_key: nil, signature_key: nil)

  # Record stuff
  @doc "Returns the encryption key of a fernet record"
  @spec encryption_key(t()) :: key()
  def encryption_key(fernet), do: fernet(fernet, :encryption_key)

  @doc "Returns the signature key of a fernet record"
  @spec signature_key(t()) :: key()
  def signature_key(fernet), do: fernet(fernet, :signature_key)

  @doc "Creates a new fernet record from a fernet token(16 bytes encryption key, 16 bytes signature key)."
  @spec new(<<_::512>>) :: t()
  def new(<<signature_key::binary-size(32), encryption_key::binary-size(32)>>) do
    fernet(encryption_key: encryption_key, signature_key: signature_key)
  end

  @doc "Checks if a signature is valid."
  @spec signature_valid?(token(), t() | key()) :: boolean()
  def signature_valid?(token, fernet) when Record.is_record(fernet, :fernet) do
    signature_valid?(token, signature_key(fernet))
  end

  def signature_valid?(token, signature_key) when is_binary(signature_key) do
    # last 32 bytes is the signature
    token_size = byte_size(token)
    <<cipher_text::binary-size(token_size - 32), received_signature::binary-size(32)>> = token
    expected_signature = Crypto.hmac(signature_key, cipher_text)

    received_signature == expected_signature
  end

  @doc "Encrypts a binary using a fernet record."
  @spec encrypt(binary(), t()) :: token()
  def encrypt(plain_text, fernet) when Record.is_record(fernet, :fernet) do
    encrypt(plain_text, encryption_key(fernet), signature_key(fernet))
  end

  @spec encrypt(binary(), key(), key()) :: token()
  def encrypt(plain_text, encryption_key, signature_key) when is_binary(plain_text) do
    # random IV
    iv = :crypto.strong_rand_bytes(16)

    # Encrypt with AES.
    cipher_text =
      :crypto.crypto_one_time(:aes_256_cbc, encryption_key, iv, plain_text,
        encrypt: true,
        padding: :pkcs_padding
      )

    # Calculate the signature and construct the final token.
    signed_data = iv <> cipher_text
    signed_data <> Crypto.hmac(signature_key, signed_data)
  end

  @doc "Decrypts a token using a fernet record."
  @spec decrypt(token(), t()) :: {:ok, binary()} | {:error, :invalid_signature}
  def decrypt(data, fernet) when Record.is_record(fernet, :fernet) do
    decrypt(data, encryption_key(fernet), signature_key(fernet))
  end

  @spec decrypt(token(), key(), key()) :: {:ok, binary()} | {:error, :invalid_signature}
  def decrypt(large_token, encryption_key, signature_key) when is_binary(large_token) do
    if signature_valid?(large_token, signature_key) do
      <<iv::binary-size(16), token::binary>> = large_token

      token_size = byte_size(token)
      <<cipher_text::binary-size(token_size - 32), _::binary-size(32)>> = token

      plain_text =
        :crypto.crypto_one_time(:aes_256_cbc, encryption_key, iv, cipher_text,
          encrypt: false,
          padding: :pkcs_padding
        )

      {:ok, plain_text}
    else
      {:error, :invalid_signature}
    end
  end
end