slip39

defmodule Slip39.Shamir do
  import Bitwise
  alias Slip39.Constants
  alias Slip39.Utils
  alias Slip39.Utils.Range
  alias Slip39.Share
  alias Slip39.Share.RawShare

  @spec interpolate(list(Share.RawShare), integer()) :: binary()
  def interpolate(raw_shares, x) do
    x_coordinates = raw_shares |> Enum.map(& &1.x) |> MapSet.new() |> MapSet.to_list()

    if length(x_coordinates) != length(raw_shares) do
      raise RuntimeError, "Invalid set of shares. Share indices must be unique."
    end

    share_value_lengths =
      raw_shares |> Enum.map(&bit_size(&1.data)) |> MapSet.new() |> MapSet.to_list()

    if length(share_value_lengths) != 1 do
      raise("Invalid set of shares. All share values must have the same length.")
    end

    found_share =
      raw_shares
      |> Enum.find(fn share ->
        share.x == x
      end)

    if is_nil(found_share) do
      # Logarithm of the product of (x_i - x) for i = 1, ... , k.
      log_prod =
        raw_shares
        |> Enum.map(&bxor(&1.x, x))
        |> Enum.map(&Constants.get!(:LOG_TABLE)[&1])
        |> Enum.sum()

      # The result accumulator is a zero filled byte array of the size
      # of the first share (because all shares have the same size)
      first_share = raw_shares |> Enum.at(0)

      result = for _ <- Range.range(first_share.data |> bit_size), do: 0

      raw_shares
      |> Enum.reduce(result, fn share, result ->
        log_basis_eval =
          Integer.mod(
            log_prod - Constants.get!(:LOG_TABLE)[bxor(share.x, x)] -
              (raw_shares
               |> Enum.map(fn other_share ->
                 bxor(share.x, other_share.x)
               end)
               |> Enum.map(&Constants.get!(:LOG_TABLE)[&1])
               |> Enum.sum()),
            255
          )

        Enum.zip(share.data |> :binary.bin_to_list(), result)
        |> Enum.map(fn {share_data, intermediate_sum} ->
          bxor(
            intermediate_sum,
            if share_data != 0 do
              Constants.get!(:EXP_TABLE)[
                (Constants.get!(:LOG_TABLE)[share_data] + log_basis_eval) |> Integer.mod(255)
              ]
            else
              0
            end
          )
        end)
      end)
      |> :binary.list_to_bin()
    else
      found_share.data
    end
  end

  def decode_mnemonics(mnemonics) do
    mnemonics
    |> Enum.map(&Share.decode!(&1))
    |> Enum.group_by(&Share.get_group_parameters(&1))
  end

  def recover_ems(groups_map) do
    if groups_map |> Map.keys() |> length() == 0 do
      raise("The set of shares is empty.")
    end

    # We fetch common_parameters from first share, since all shares
    # must have the same common parameters.
    common_params =
      groups_map |> Map.values() |> Enum.at(0) |> Enum.at(0) |> Share.get_common_parameters()

    if groups_map |> Map.keys() |> length() < common_params.group_threshold do
      raise(
        "Insufficient number of mnemonic groups. The required number of groups is #{common_params.group_threshold}."
      )
    end

    if groups_map |> Map.keys() |> length() != common_params.group_threshold do
      raise(
        "Wrong number of mnemonic groups. Expected #{common_params.group_threshold}" <>
          "groups, but #{groups_map |> Map.keys() |> length()} were provided."
      )
    end

    # Let's check that each group have at least group_threshold count of share
    for group_parameters <- groups_map |> Map.keys() do
      if length(groups_map[group_parameters]) < group_parameters.member_threshold do
        # Let's take the first share to deduce the group prefix words
        first_share = groups_map[group_parameters] |> Enum.at(0)
        words = first_share |> Share.encode!()

        prefix =
          words |> Enum.slice(0, Constants.get!(:GROUP_PREFIX_LENGTH_WORDS)) |> Enum.join(" ")

        raise "Wrong number of mnemonics. " <>
                "Expected at least#{group_parameters.member_threshold} mnemonics starting with " <>
                "words '#{prefix}', but #{length(groups_map[group_parameters])} were provided."
      end
    end

    group_raw_shares =
      groups_map
      |> Map.keys()
      |> Enum.map(fn params ->
        groups_map[params] |> Slip39.Share.get_raw_shares()

        %Share.RawShare{
          x: params.group_index,
          data:
            recover_secret(
              params.member_threshold,
              groups_map[params] |> Slip39.Share.get_raw_shares()
            )
        }
      end)

    {recover_secret(common_params.group_threshold, group_raw_shares), common_params}
  end

  def recover_secret(1, raw_shares) do
    first_raw_share = raw_shares |> Enum.at(0)
    first_raw_share.data
  end

  def recover_secret(_threshold, raw_shares) do
    shared_secret = interpolate(raw_shares, Constants.get!(:SECRET_INDEX))
    digest_share = interpolate(raw_shares, Constants.get!(:DIGEST_INDEX))

    digest_size = Constants.get!(:DIGEST_LENGTH_BYTES)

    <<digest::binary-size(digest_size), random_part::binary>> =
      digest_share

    if digest != create_digest(random_part, shared_secret) do
      raise "Invalid digest of the shared secret."
    end

    shared_secret
  end

  @spec create_digest(binary(), binary()) :: binary()
  def create_digest(random_data, shared_secret) do
    digest_size = Constants.get!(:DIGEST_LENGTH_BYTES)

    <<digest::binary-size(digest_size), _rest::binary>> =
      :crypto.mac(:hmac, :sha256, random_data, shared_secret)

    <<digest::binary-size(digest_size)>>
  end

  @spec split_ems!(integer(), list(), binary(), any(), any()) :: list()
  def split_ems!(group_threshold, groups, encrypted_master_secret, iteration_exponent, identifier) do
    # todo : validation

    group_shares =
      split_secret!(
        group_threshold,
        length(groups),
        encrypted_master_secret
      )

    for {{member_threshold, member_count}, %RawShare{x: group_index, data: group_secret}} <-
          Enum.zip(groups, group_shares) do
      for %RawShare{x: member_index, data: value} <-
            split_secret!(member_threshold, member_count, group_secret) do
        %Share{
          identifier: identifier,
          iteration_exponent: iteration_exponent,
          group_index: group_index,
          group_threshold: group_threshold,
          group_count: length(groups),
          member_index: member_index,
          member_threshold: member_threshold,
          data: value
        }
      end
    end
  end

  @spec split_secret!(integer(), integer(), binary()) :: list(%RawShare{})
  def split_secret!(threshold, share_count, shared_secret) do
    if threshold < 1,
      do: raise("The requested threshold must be a positive integer.")

    if threshold > share_count,
      do: raise("The requested threshold must not exceed the number of shares.")

    if share_count > Constants.get!(:MAX_SHARE_COUNT),
      do:
        raise(
          RuntimeError,
          "The requested number of shares must not exceed #{Constants.get!(:MAX_SHARE_COUNT)}."
        )

    if threshold == 1 do
      # If the threshold is 1, then the digest of the shared secret is not used.
      for i <- Range.range(share_count), do: %RawShare{x: i, data: shared_secret}
    else
      random_share_count = threshold - 2

      shares =
        for i <- Range.range(random_share_count),
            do: %RawShare{x: i, data: Utils.rnd_bytes(byte_size(shared_secret))}

      random_part =
        Utils.rnd_bytes(byte_size(shared_secret) - Constants.get!(:DIGEST_LENGTH_BYTES))

      digest = create_digest(random_part, shared_secret)

      base_shares =
        shares ++
          [
            %RawShare{x: Constants.get!(:DIGEST_INDEX), data: digest <> random_part},
            %RawShare{x: Constants.get!(:SECRET_INDEX), data: shared_secret}
          ]

      shares ++
        for i <- Range.range(random_share_count, share_count),
            do: %RawShare{x: i, data: interpolate(base_shares, i)}
    end
  end

  @spec random_identifier() :: non_neg_integer()
  def random_identifier() do
    id_length_bits = Constants.get!(:ID_LENGTH_BITS)
    byte_count = Utils.bits_to_bytes(id_length_bits)
    <<identifier::size(id_length_bits), _rest::bitstring>> = Utils.rnd_bytes(byte_count)

    identifier
  end

  @spec compute_exp_log_tables() :: {%{integer() => integer()}, %{integer() => integer()}}
  def compute_exp_log_tables do
    # This code is not very efficient, but since it is called
    # for module attribute initialization, it will be run twice
    # (once for each attribute) at compile time, I won't spend too
    # much time changing this for now.
    # Anyway, any PR is still welcome.

    exp = for x <- Range.range(255), do: x
    log = for x <- Range.range(256), do: x
    poly = 1

    {_poly, exp, log} =
      Range.range(255)
      |> Enum.reduce({poly, exp, log}, fn i, {poly, exp, log} ->
        exp = exp |> List.replace_at(i, poly)
        log = log |> List.replace_at(poly, i)

        poly = bxor(poly <<< 1, poly)

        poly =
          if (poly &&& 0x100) > 0 do
            poly |> bxor(0x11B)
          else
            poly
          end

        {poly, exp, log}
      end)

    {exp |> Utils.enum_to_indexed_map(), log |> Utils.enum_to_indexed_map()}
  end
end