ex_web3

defmodule ExWeb3 do
  @moduledoc """
  Helpers for common operations around EVM networks.
  """

  import Bitwise

  require Logger

  @type keccak_hash :: binary()

  @ethereum_signature_prefix "\x19Ethereum Signed Message:\n"

  @doc """
  Simply returns the rightmost n bits of a binary.

  ## Examples

      iex> ExWeb3.mask(0b111101111, 6)
      0b101111
      iex> ExWeb3.mask(0b101101, 3)
      0b101
      iex> ExWeb3.mask(0b011, 1)
      1
      iex> ExWeb3.mask(0b011, 0)
      0
      iex> ExWeb3.mask(0b010, 1)
      0
      iex> ExWeb3.mask(0b010, 20)
      0b010
  """
  @spec mask(integer(), integer()) :: integer()
  def mask(n, bits) when is_integer(n) do
    # Calculates n bitwise-and 0b1111...<bits>
    n &&& (2 <<< (bits - 1)) - 1
  end

  @doc """
  Simply returns the rightmost n bits of a binary.

  ## Examples

      iex> ExWeb3.mask_bitstring(<<0b111101111>>, 6)
      <<0b101111::size(6)>>
      iex> ExWeb3.mask_bitstring(<<0b101101>>, 3)
      <<0b101::size(3)>>
      iex> ExWeb3.mask_bitstring(<<0b011>>, 1)
      <<1::size(1)>>
      iex> ExWeb3.mask_bitstring(<<0b011>>, 0)
      <<>>
      iex> ExWeb3.mask_bitstring(<<0b010>>, 1)
      <<0::size(1)>>
      iex> ExWeb3.mask_bitstring(<<0b010>>, 20)
      <<0, 0, 2::size(4)>>
      iex> ExWeb3.mask_bitstring(<<0b010>>, 3)
      <<0b010::size(3)>>
      iex> ExWeb3.mask_bitstring(<<>>, 3)
      <<0b000::size(3)>>
  """
  @spec mask_bitstring(bitstring(), integer()) :: bitstring()
  def mask_bitstring(b, bits) do
    size = bit_size(b)
    skip_size = max(size - bits, 0)
    padding = max(bits - size, 0)

    <<_::size(skip_size), included_part::bits>> = b

    <<0::size(padding), included_part::bitstring>>
  end

  @doc """
  Returns the keccak sha256 of a given input.

  ## Examples

      iex> ExWeb3.kec("hello world")
      <<71, 23, 50, 133, 168, 215, 52, 30, 94, 151, 47, 198, 119, 40, 99,
             132, 248, 2, 248, 239, 66, 165, 236, 95, 3, 187, 250, 37, 76, 176,
             31, 173>>
      iex> ExWeb3.kec(<<0x01, 0x02, 0x03>>)
      <<241, 136, 94, 218, 84, 183, 160, 83, 49, 140, 212, 30, 32, 147, 34,
             13, 171, 21, 214, 83, 129, 177, 21, 122, 54, 51, 168, 59, 253, 92,
             146, 57>>
  """
  @spec kec(binary()) :: keccak_hash
  def kec(data) do
    ExKeccak.hash_256(data)
  end

  @doc """
  Returned a binary padded to given length in bytes. Fails if
  binary is longer than desired length.

  ## Examples

      iex> ExWeb3.pad(<<9>>, 5)
      <<0, 0, 0, 0, 9>>
      iex> ExWeb3.pad(<<9>>, 1)
      <<9>>
      iex> ExWeb3.pad(<<9, 9>>, 1)
      ** (RuntimeError) Binary too long for padding
  """
  @spec pad(binary(), integer()) :: binary()
  def pad(binary, desired_length) do
    desired_bits = desired_length * 8

    case byte_size(binary) do
      0 ->
        <<0::size(desired_bits)>>

      x when x <= desired_length ->
        padding_bits = (desired_length - x) * 8
        <<0::size(padding_bits)>> <> binary

      _ ->
        raise "Binary too long for padding"
    end
  end

  @doc """
  Similar to `:binary.encode_unsigned/1`, except we encode `0` as
  `<<>>`, the empty string. This is because the specification says that
  we cannot have any leading zeros, and so having <<0>> by itself is
  leading with a zero and prohibited.

  ## Examples

      iex> ExWeb3.encode_unsigned(0)
      <<>>
      iex> ExWeb3.encode_unsigned(5)
      <<5>>
      iex> ExWeb3.encode_unsigned(5_000_000)
      <<76, 75, 64>>
  """
  @spec encode_unsigned(non_neg_integer()) :: binary()
  def encode_unsigned(0), do: <<>>
  def encode_unsigned(n), do: :binary.encode_unsigned(n)

  @doc """
  Similar to `:binary.decode_unsigned/1`, except we decode `<<>>` back to `0`,
  which is a common practice in Ethereum, since we cannot have **any** leading
  zeros.

  ## Examples

      iex> ExWeb3.decode_unsigned(<<>>)
      0
      iex> ExWeb3.decode_unsigned(<<5>>)
      5
      iex> ExWeb3.decode_unsigned(<<76, 75, 64>>)
      5_000_000
  """
  @spec decode_unsigned(binary()) :: non_neg_integer()
  def decode_unsigned(<<>>), do: 0
  def decode_unsigned(bin), do: :binary.decode_unsigned(bin)

  @doc """
  Convert hex encoded string to integer.

  ## Examples

      iex> ExWeb3.hex_to_int!("0x1")
      1
      iex> ExWeb3.hex_to_int!("0xA")
      10
  """
  @spec hex_to_int!(String.t()) :: integer() | no_return()
  def hex_to_int!("0x" <> hex) do
    case Integer.parse(hex, 16) do
      {int, ""} ->
        int

      _ ->
        raise "invalid hex string"
    end
  end

  @doc """
  Simple wrapper for decoding hex data.

  ## Examples

      iex> ExWeb3.from_hex("aabbcc")
      <<0xaa, 0xbb, 0xcc>>
  """
  @spec from_hex(String.t()) :: binary()
  def from_hex(hex_data), do: Base.decode16!(hex_data, case: :lower)

  @doc """
  Simple wrapper to generate hex.

  ## Examples

      iex> ExWeb3.to_hex(<<0xaa, 0xbb, 0xcc>>)
      "aabbcc"
  """
  @spec to_hex(binary()) :: String.t()
  def to_hex(bin), do: Base.encode16(bin, case: :lower)

  @doc """
  Simple wrapper for decoding hex data.

  ## Examples

      iex> ExWeb3.from_prefixed_hex!("0xaabbcc")
      <<0xaa, 0xbb, 0xcc>>
  """
  @spec from_prefixed_hex!(String.t()) :: binary()
  def from_prefixed_hex!("0x" <> hex_data), do: Base.decode16!(hex_data, case: :mixed)

  @doc """
  Simple wrapper for decoding hex data.

  ## Examples

      iex> ExWeb3.from_prefixed_hex("0xaabbcc")
      {:ok, <<0xaa, 0xbb, 0xcc>>}
  """
  @spec from_prefixed_hex(String.t()) :: {:ok, binary()} | {:error, :invalid_hex}
  def from_prefixed_hex("0x" <> hex_data) do
    case Base.decode16(hex_data, case: :mixed) do
      {:ok, decoded} ->
        {:ok, decoded}

      :error ->
        {:error, :invalid_hex}
    end
  end

  @doc """
  Encoder from binary to hex data.

  ## Examples

      iex> ExWeb3.to_prefixed_hex(<<0xaa, 0xbb, 0xcc>>)
      "0xaabbcc"
  """
  @spec to_prefixed_hex(binary()) :: String.t()
  def to_prefixed_hex(bin), do: "0x" <> to_hex(bin)

  @doc """
  Simple wrapper for decoding ABI encoded hex data.
  """
  @spec from_abi(String.t()) :: {integer(), <<_::256>>, <<_::256>>} | no_return
  def from_abi("0x" <> hex_data) do
    [{v, r, s}] = ABI.decode("(uint8,bytes32,bytes32)", Base.decode16!(hex_data, case: :mixed))
    {v, r, s}
  end

  @doc """
  Get public key from private key.
  """
  @spec get_public_key(binary()) :: binary()
  def get_public_key(<<private_key::binary-size(32)>>) do
    {:ok, public_key} = ExSecp256k1.create_public_key(private_key)
    public_key
  end

  def get_public_key(<<encoded_private_key::binary-size(64)>>) do
    private_key = Base.decode16!(encoded_private_key, case: :mixed)
    {:ok, public_key} = ExSecp256k1.create_public_key(private_key)
    public_key
  end

  @doc """
  Get Ethereum address from private or public key.
  """
  @spec get_address(binary()) :: String.t()
  def get_address(<<private_key::binary-size(32)>>) do
    <<4::size(8), key::binary-size(64)>> = private_key |> get_public_key()
    <<_::binary-size(12), eth_address::binary-size(20)>> = kec(key)
    "0x#{Base.encode16(eth_address, case: :lower)}"
  end

  def get_address(<<encoded_private_key::binary-size(64)>>) do
    public_key = Base.decode16!(encoded_private_key, case: :mixed) |> get_public_key()
    <<4::size(8), key::binary-size(64)>> = public_key
    <<_::binary-size(12), eth_address::binary-size(20)>> = kec(key)
    "0x#{Base.encode16(eth_address, case: :lower)}"
  end

  def get_address(<<4::size(8), key::binary-size(64)>>) do
    <<_::binary-size(12), eth_address::binary-size(20)>> = kec(key)
    "0x#{Base.encode16(eth_address, case: :lower)}"
  end

  def get_address(<<encoded_public_key::binary-size(130)>>) do
    <<4::size(8), key::binary-size(64)>> = Base.decode16!(encoded_public_key, case: :mixed)
    <<_::binary-size(12), eth_address::binary-size(20)>> = kec(key)
    "0x#{Base.encode16(eth_address, case: :lower)}"
  end

  @doc """
  Generate random expirable challenge for signature.
  """
  @spec gen_eth_challenge(String.t(), integer()) :: {String.t(), String.t(), integer()}
  def gen_eth_challenge(intention, challenge_expiration \\ 300) do
    challenge = to_hex(:crypto.strong_rand_bytes(8))
    expiration = :os.system_time(:second) + challenge_expiration
    login_challenge = "#{intention}:#{challenge}:expires:#{expiration}"
    {login_challenge, challenge, expiration}
  end

  @doc """
  Check if Ethereum login challenge has expired.
  """
  @spec is_eth_signed_challenge_expired?(String.t()) :: boolean()
  def is_eth_signed_challenge_expired?(
        <<"login:", _::binary-size(16), ":expires:", expires::binary>>
      ) do
    case Integer.parse(expires) do
      {expires_at, _} ->
        :os.system_time(:second) > expires_at

      _ ->
        false
    end
  end

  def is_eth_signed_challenge_expired?(_), do: false

  @doc """
  Validate if Ethereum challenge digest was signed with a proper wallet.
  """
  @spec is_eth_signed_challenge_valid?(String.t(), String.t(), String.t()) :: boolean()
  def is_eth_signed_challenge_valid?(wallet, challenge, raw_signature) do
    {signature, v} =
      case raw_signature do
        <<_::binary-size(194)>> ->
          {v, r, s} = from_abi(raw_signature)
          {r <> s, v}

        _ ->
          <<signature::binary-size(64), v::integer-size(8)>> = from_prefixed_hex!(raw_signature)
          {signature, v}
      end

    hash = eth_hash_message!(challenge)

    case recover(hash, signature, get_secp256k1_recovery_from_eth_signature_recovery(v)) do
      {:ok, public_key} ->
        get_address(public_key) === wallet and
          verify(hash, signature, public_key)

      _ ->
        false
    end
  end

  @doc """
  Verifies a that a given signature for a given digest is valid against a public key.
  """
  @spec verify(binary(), binary(), binary()) :: boolean()
  def verify(digest, signature, public_key) do
    case ExSecp256k1.verify(digest, signature, public_key) do
      :ok -> true
      {:error, _error} -> false
    end
  end

  @doc """
  Recovers a public key from a given signature structured as a hex string returned by ethers.js e.g:

      String.length("0x" <> signature <> recovery_id) == 67
  """
  @spec eth_recover(binary(), String.t()) :: {:ok, binary()} | {:error, any()}
  def eth_recover(digest, "0x" <> _ = hex_signature) do
    with <<signature::binary-size(64), v::integer-size(8)>> <- from_prefixed_hex(hex_signature) do
      recovery_id = get_secp256k1_recovery_from_eth_signature_recovery(v)
      recover(digest, signature, recovery_id)
    end
  end

  @doc """
  Recovers a public key from a given signature for a given digest.
  Note, the key is returned in DER format (with a leading 0x04 to indicate
  that it's an octet-string).
  """
  @spec recover(binary(), binary(), integer()) :: {:ok, binary()} | {:error, String.t()}
  def recover(digest, signature, recovery_id) do
    case ExSecp256k1.recover_compact(digest, signature, recovery_id) do
      {:ok, public_key} -> {:ok, public_key}
      {:error, reason} -> {:error, to_string(reason)}
    end
  end

  @doc """
  Signs message with private key, returns signature parts {r, s, v}
  """
  @spec sign(binary, binary) :: {:error, atom} | {:ok, {binary, binary, non_neg_integer}}
  def sign(message, private_key) do
    case ExSecp256k1.sign(message, private_key) do
      {:ok, signature} -> {:ok, signature}
      {:error, error} -> {:error, error}
    end
  end

  @spec eth_sign(binary, binary) :: {:error, atom} | {:ok, {binary, binary, non_neg_integer}}
  def eth_sign(message, private_key) do
    with hash <- eth_hash_message!(message),
         {:ok, {r, s, v}} <- sign(hash, private_key) do
      {:ok, {r, s, get_eth_signature_recovery_from_secp256k1_recovery(v)}}
    end
  end

  @doc """
  Create digest for Ethereum signature for a generic binary data.
  """
  @spec eth_hash_message!(binary()) :: keccak_hash | no_return()
  def eth_hash_message!(message) do
    kec("#{@ethereum_signature_prefix}#{byte_size(message)}#{message}")
  end

  # Secp256k1 signature consists of 3 parts VRS, where V is recovery_id
  # for security reasons ethereum modifies this recovery_id.
  # In order to prepare signature to be useed with eth chain it is needed to modify
  # it in the same way
  # @link https://github.com/ethereum/EIPs/blob/master/EIPS/eip-155.md
  # @link https://ethereum.stackexchange.com/questions/42455/during-ecdsa-signing-how-do-i-generate-the-recovery-id
  # @link https://ethereum.stackexchange.com/questions/15766/what-does-v-r-s-in-eth-gettransactionbyhash-mean
  defp get_secp256k1_recovery_from_eth_signature_recovery(recovery_id) do
    unless recovery_id in [27, 28] do
      Logger.warn("[#{__MODULE__}] unexpected eth recovery id")
    end

    recovery_id - 27
  end

  # see `get_secp256k1_recovery_from_eth_signature_recovery/1` docs
  defp get_eth_signature_recovery_from_secp256k1_recovery(recovery_id) do
    unless recovery_id in [0, 1] do
      Logger.warn("[#{__MODULE__}] unexpected secp256k1 recovery id")
    end

    recovery_id + 27
  end
end