defmodule AirPlay.V2.Srp do
  @moduledoc """
  AirPlay 2 SRP-6a client proof generation.

  AirPlay 2 transient pairing uses username `Pair-Setup`, password `3939`, the
  RFC 5054 3072-bit group, generator 5 and SHA-512.
  """

  import Bitwise, only: [bxor: 2]

  defstruct [:client_public_key, :client_pk, :client_proof, :session_key, :server_proof]

  @type t :: %__MODULE__{
          client_public_key: binary(),
          client_pk: binary(),
          client_proof: binary(),
          session_key: binary(),
          server_proof: binary()
        }

  @n_hex """
  FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E08
  8A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B
  302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9
  A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE6
  49286651ECE45B3DC2007CB8A163BF0598DA48361C55D39A69163FA8
  FD24CF5F83655D23DCA3AD961C62F356208552BB9ED529077096966D
  670C354E4ABC9804F1746C08CA18217C32905E462E36CE3BE39E772C
  180E86039B2783A2EC07A28FB5C55DF06F4C52C9DE2BCBF695581718
  3995497CEA956AE515D2261898FA051015728E5A8AAAC42DAD33170D
  04507A33A85521ABDF1CBA64ECFB850458DBEF0A8AEA71575D060C7D
  B3970F85A6E1E4C7ABF5AE8CDB0933D71E8C94E04A25619DCEE3D226
  1AD2EE6BF12FFA06D98A0864D87602733EC86A64521F2B18177B200C
  BBE117577A615D6C770988C0BAD946E208E24FA074E5AB3143DB5BFCE0FD108E
  4B82D120A93AD2CAFFFFFFFFFFFFFFFF
  """

  @n_bin @n_hex |> String.replace(~r/\s+/, "") |> Base.decode16!()
  @n :binary.decode_unsigned(@n_bin)
  @n_len byte_size(@n_bin)
  @g 5

  @doc "Create an AirPlay 2 SRP client proof."
  @spec create(binary(), binary(), binary(), binary(), keyword()) :: t()
  def create(username, password, server_public_key, salt, opts \\ []) do
    private = Keyword.get_lazy(opts, :private_key, fn -> :crypto.strong_rand_bytes(32) end)
    a = decode(private)
    big_b = decode(server_public_key)
    big_a = mod_pow(@g, a, @n)

    k = hash_int(pad(@n) <> pad(@g))
    x = hash_int(salt <> hash(username <> ":" <> password))
    u = hash_int(pad(big_a) <> pad(big_b))
    gx = mod_pow(@g, x, @n)
    base = positive_mod(big_b - positive_mod(k * gx, @n), @n)
    exponent = a + u * x
    shared_secret = mod_pow(base, exponent, @n)
    session_key = hash(pad(shared_secret))

    proof =
      hash(
        xor_bytes(hash(pad(@n)), hash(:binary.encode_unsigned(@g))) <>
          hash(username) <>
          salt <>
          pad(big_a) <>
          pad(big_b) <>
          session_key
      )

    %__MODULE__{
      client_public_key: pad(big_a),
      client_pk: pad(big_a),
      client_proof: proof,
      session_key: session_key,
      server_proof: hash(pad(big_a) <> proof <> session_key)
    }
  end

  @doc "Convenience helper for transient AP2 pairing."
  @spec transient(binary(), binary(), keyword()) :: t()
  def transient(server_public_key, salt, opts \\ []) do
    create("Pair-Setup", "3939", server_public_key, salt, opts)
  end

  @doc "RFC 5054 3072-bit group modulus."
  @spec n() :: pos_integer()
  def n, do: @n

  @doc "RFC 5054 group generator."
  @spec g() :: pos_integer()
  def g, do: @g

  @doc "Compatibility wrapper for SRP session generation."
  @spec session(binary(), binary(), binary(), binary(), keyword()) :: t()
  def session(username, password, server_public_key, salt, opts \\ []) do
    opts =
      case Keyword.fetch(opts, :a) do
        {:ok, a} -> Keyword.put(opts, :private_key, :binary.encode_unsigned(a))
        :error -> opts
      end

    create(username, password, server_public_key, salt, opts)
  end

  defp hash(data), do: :crypto.hash(:sha512, data)
  defp hash_int(data), do: data |> hash() |> decode()
  defp decode(data), do: :binary.decode_unsigned(data)

  defp pad(int) do
    int
    |> :binary.encode_unsigned()
    |> then(fn bin -> :binary.copy(<<0>>, max(@n_len - byte_size(bin), 0)) <> bin end)
  end

  defp mod_pow(base, exponent, modulus) do
    base
    |> :binary.encode_unsigned()
    |> :crypto.mod_pow(:binary.encode_unsigned(exponent), :binary.encode_unsigned(modulus))
    |> decode()
  end

  defp positive_mod(value, modulus),
    do: value |> rem(modulus) |> Kernel.+(modulus) |> rem(modulus)

  defp xor_bytes(left, right) do
    left_bytes = :binary.bin_to_list(left)
    right_bytes = :binary.bin_to_list(right)

    left_bytes
    |> Enum.zip(right_bytes)
    |> Enum.map(fn {l, r} -> bxor(l, r) end)
    |> :binary.list_to_bin()
  end
end