extra/lib/plausible/auth/sso.ex

defmodule Plausible.Auth.SSO do
  @moduledoc """
  API for SSO.
  """

  import Ecto.Changeset
  import Ecto.Query

  alias Plausible.Auth
  alias Plausible.Auth.SSO
  alias Plausible.Billing.Subscription
  alias Plausible.Repo
  alias Plausible.Teams

  use Plausible.Auth.SSO.Domain.Status

  require Plausible.Billing.Subscription.Status

  @type policy_attr() ::
          {:sso_default_role, Teams.Policy.sso_member_role()}
          | {:sso_session_timeout_minutes, non_neg_integer()}

  @spec get_integration_for(Teams.Team.t()) :: {:ok, SSO.Integration.t()} | {:error, :not_found}
  def get_integration_for(%Teams.Team{} = team) do
    query = integration_query() |> where([i], i.team_id == ^team.id)

    if integration = Repo.one(query) do
      {:ok, integration}
    else
      {:error, :not_found}
    end
  end

  @spec get_integration(String.t()) :: {:ok, SSO.Integration.t()} | {:error, :not_found}
  def get_integration(identifier) when is_binary(identifier) do
    query = integration_query() |> where([i], i.identifier == ^identifier)

    if integration = Repo.one(query) do
      {:ok, integration}
    else
      {:error, :not_found}
    end
  end

  defp integration_query() do
    from(i in SSO.Integration,
      inner_join: t in assoc(i, :team),
      as: :team,
      left_join: d in assoc(i, :sso_domains),
      as: :sso_domains,
      preload: [team: t, sso_domains: d]
    )
  end

  @spec initiate_saml_integration(Teams.Team.t()) :: SSO.Integration.t()
  def initiate_saml_integration(team) do
    changeset = SSO.Integration.init_changeset(team)

    Repo.insert_with_audit!(
      changeset,
      "saml_integration_initiated",
      %{team_id: team.id},
      on_conflict: [set: [updated_at: NaiveDateTime.utc_now(:second)]],
      conflict_target: :team_id,
      returning: true
    )
  end

  @spec update_integration(SSO.Integration.t(), map()) ::
          {:ok, SSO.Integration.t()} | {:error, Ecto.Changeset.t()}
  def update_integration(integration, params) do
    changeset = SSO.Integration.update_changeset(integration, params)

    case Repo.update_with_audit(changeset, "sso_integration_updated", %{
           team_id: integration.team_id
         }) do
      {:ok, integration} -> {:ok, integration}
      {:error, changeset} -> {:error, changeset.changes.config}
    end
  end

  @spec provision_user(SSO.Identity.t()) ::
          {:ok, :standard | :sso | :integration, Teams.Team.t(), Auth.User.t()}
          | {:error, :integration_not_found | :over_limit}
          | {:error, :multiple_memberships | :active_personal_team, Teams.Team.t(), Auth.User.t()}
  def provision_user(identity) do
    case find_user(identity) do
      {:ok, :standard, user, integration, domain} ->
        provision_standard_user(user, identity, integration, domain)

      {:ok, :sso, user, integration, domain} ->
        provision_sso_user(user, identity, integration, domain)

      {:ok, :integration, _, integration, domain} ->
        provision_identity(identity, integration, domain)

      {:error, :not_found} ->
        {:error, :integration_not_found}
    end
  end

  @spec deprovision_user!(Auth.User.t()) :: Auth.User.t()
  def deprovision_user!(%{type: :standard} = user), do: user

  def deprovision_user!(user) do
    user = Repo.preload(user, [:sso_integration, :sso_domain])

    :ok = Auth.UserSessions.revoke_all(user)

    user
    |> Ecto.Changeset.change()
    |> Ecto.Changeset.put_change(:type, :standard)
    |> Ecto.Changeset.put_change(:sso_identity_id, nil)
    |> Ecto.Changeset.put_assoc(:sso_integration, nil)
    |> Ecto.Changeset.put_assoc(:sso_domain, nil)
    |> Repo.update_with_audit!("sso_user_deprovioned", %{team_id: user.sso_integration.team_id})
  end

  @spec update_policy(Teams.Team.t(), [policy_attr()]) ::
          {:ok, Teams.Team.t()} | {:error, Ecto.Changeset.t()}
  def update_policy(team, attrs \\ []) do
    params = Map.new(attrs)

    changeset =
      team
      |> Ecto.Changeset.cast(%{policy: params}, [])
      |> Ecto.Changeset.cast_embed(:policy, with: &Teams.Policy.update_changeset/2)

    case Repo.update_with_audit(changeset, "sso_policy_updated", %{team_id: team.id}) do
      {:ok, integration} -> {:ok, integration}
      {:error, changeset} -> {:error, changeset.changes.policy}
    end
  end

  @spec set_force_sso(Teams.Team.t(), Teams.Policy.force_sso_mode()) ::
          {:ok, Teams.Team.t()}
          | {:error,
             :no_integration
             | :no_domain
             | :no_verified_domain
             | :owner_2fa_disabled
             | :no_sso_user}
  def set_force_sso(team, mode) do
    with :ok <- check_force_sso(team, mode) do
      params = %{policy: %{force_sso: mode}}

      team
      |> Ecto.Changeset.cast(params, [])
      |> Ecto.Changeset.cast_embed(:policy,
        with: &Teams.Policy.force_sso_changeset(&1, &2.force_sso)
      )
      |> Repo.update_with_audit("sso_force_mode_changed", %{team_id: team.id})
    end
  end

  @spec check_force_sso(Teams.Team.t(), Teams.Policy.force_sso_mode()) ::
          :ok
          | {:error,
             :no_integration
             | :no_domain
             | :no_verified_domain
             | :owner_2fa_disabled
             | :no_sso_user}
  def check_force_sso(_team, :none), do: :ok

  def check_force_sso(team, :all_but_owners) do
    with :ok <- check_integration_configured(team),
         :ok <- check_sso_user_present(team) do
      check_owners_2fa_enabled(team)
    end
  end

  @spec check_can_remove_integration(SSO.Integration.t()) ::
          :ok | {:error, :force_sso_enabled | :sso_users_present}
  def check_can_remove_integration(integration) do
    team = Repo.preload(integration, :team).team

    cond do
      team.policy.force_sso != :none ->
        {:error, :force_sso_enabled}

      check_sso_user_present(integration) == :ok ->
        {:error, :sso_users_present}

      true ->
        :ok
    end
  end

  @spec remove_integration(SSO.Integration.t(), Keyword.t()) ::
          :ok | {:error, :force_sso_enabled | :sso_users_present}
  def remove_integration(integration, opts \\ []) do
    force_deprovision? = Keyword.get(opts, :force_deprovision?, false)
    check = check_can_remove_integration(integration)

    case {check, force_deprovision?} do
      {:ok, _} ->
        {:ok, :ok} =
          Repo.transaction(fn ->
            integration = Repo.preload(integration, :sso_domains)
            Enum.each(integration.sso_domains, &SSO.Domains.cancel_verification(&1.domain))

            Repo.delete_with_audit!(integration, "sso_integration_removed", %{
              team_id: integration.team_id
            })

            :ok
          end)

        :ok

      {{:error, :sso_users_present}, true} ->
        {:ok, :ok} =
          Repo.transaction(fn ->
            users = Repo.preload(integration, :users).users
            integration = Repo.preload(integration, :sso_domains)
            Enum.each(users, &deprovision_user!/1)
            Enum.each(integration.sso_domains, &SSO.Domains.cancel_verification(&1.domain))

            Repo.delete_with_audit!(integration, "sso_integration_removed", %{
              team_id: integration.team_id
            })

            :ok
          end)

        :ok

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

  @spec check_ready_to_provision(Auth.User.t(), Teams.Team.t()) ::
          :ok | {:error, :not_a_member | :multiple_memberships | :active_personal_team}
  def check_ready_to_provision(%{type: :sso} = _user, _team), do: :ok

  def check_ready_to_provision(user, team) do
    result =
      with :ok <- ensure_team_member(team, user),
           :ok <- ensure_one_membership(user, team) do
        ensure_empty_personal_team(user, team)
      end

    case result do
      :ok -> :ok
      {:error, :integration_not_found} -> {:error, :not_a_member}
      {:error, :multiple_memberships, _, _} -> {:error, :multiple_memberships}
      {:error, :active_personal_team, _, _} -> {:error, :active_personal_team}
    end
  end

  defp check_integration_configured(team) do
    integrations =
      Repo.all(
        from(
          i in SSO.Integration,
          left_join: d in assoc(i, :sso_domains),
          where: i.team_id == ^team.id,
          preload: [sso_domains: d]
        )
      )

    domains = Enum.flat_map(integrations, & &1.sso_domains)
    no_verified_domains? = Enum.all?(domains, &(&1.status != Status.verified()))

    cond do
      integrations == [] -> {:error, :no_integration}
      domains == [] -> {:error, :no_domain}
      no_verified_domains? -> {:error, :no_verified_domain}
      true -> :ok
    end
  end

  defp check_sso_user_present(%Teams.Team{} = team) do
    sso_user_count =
      Repo.aggregate(
        from(
          tm in Teams.Membership,
          inner_join: u in assoc(tm, :user),
          where: tm.team_id == ^team.id,
          where: tm.role != :guest,
          where: u.type == :sso
        ),
        :count
      )

    if sso_user_count > 0 do
      :ok
    else
      {:error, :no_sso_user}
    end
  end

  defp check_sso_user_present(%SSO.Integration{} = integration) do
    sso_user_count =
      Repo.aggregate(
        from(
          i in SSO.Integration,
          inner_join: u in assoc(i, :users),
          inner_join: tm in assoc(u, :team_memberships),
          on: tm.team_id == i.team_id,
          where: i.id == ^integration.id,
          where: tm.role != :guest,
          where: u.type == :sso
        ),
        :count
      )

    if sso_user_count > 0 do
      :ok
    else
      {:error, :no_sso_user}
    end
  end

  defp check_owners_2fa_enabled(team) do
    disabled_2fa_count =
      Repo.aggregate(
        from(
          tm in Teams.Membership,
          inner_join: u in assoc(tm, :user),
          where: tm.team_id == ^team.id,
          where: tm.role == :owner,
          where: u.totp_enabled == false or is_nil(u.totp_secret)
        ),
        :count
      )

    if disabled_2fa_count == 0 do
      :ok
    else
      {:error, :owner_2fa_disabled}
    end
  end

  defp find_user(identity) do
    case find_user_with_fallback(identity) do
      {:ok, type, user, integration, domain} ->
        {:ok, type, Repo.preload(user, [:sso_integration, :sso_domain]), integration, domain}

      {:error, _} = error ->
        error
    end
  end

  defp find_user_with_fallback(identity) do
    with {:ok, integration} <- get_integration(identity.integration_id),
         {:error, :not_found} <- find_by_identity(identity, integration) do
      find_by_email(identity.email, integration)
    end
  end

  defp find_by_identity(identity, integration) do
    if user =
         Repo.get_by(Auth.User,
           sso_integration_id: integration.id,
           sso_identity_id: identity.id,
           type: :sso
         ) do
      with {:ok, sso_domain} <- SSO.Domains.lookup(identity.email),
           :ok <- check_domain_integration_match(sso_domain, integration) do
        user = Repo.preload(user, sso_integration: :team)

        {:ok, user.type, user, user.sso_integration, sso_domain}
      end
    else
      {:error, :not_found}
    end
  end

  defp find_by_email(email, integration) do
    with {:ok, sso_domain} <- SSO.Domains.lookup(email),
         :ok <- check_domain_integration_match(sso_domain, integration) do
      case find_in_team_by_email(sso_domain.sso_integration.team, email) do
        {:ok, user} ->
          {:ok, user.type, user, sso_domain.sso_integration, sso_domain}

        {:error, :not_found} ->
          {:ok, :integration, nil, sso_domain.sso_integration, sso_domain}
      end
    end
  end

  defp find_in_team_by_email(team, email) do
    result =
      Repo.one(
        from(
          u in Auth.User,
          inner_join: tm in assoc(u, :team_memberships),
          where: u.email == ^email,
          where: tm.team_id == ^team.id,
          where: tm.role != :guest
        )
      )

    if result do
      {:ok, result}
    else
      {:error, :not_found}
    end
  end

  defp check_domain_integration_match(sso_domain, integration) do
    if sso_domain.sso_integration_id == integration.id do
      :ok
    else
      {:error, :not_found}
    end
  end

  defp provision_sso_user(user, identity, integration, domain) do
    changeset =
      user
      |> change()
      |> put_change(:email, identity.email)
      |> put_change(:name, identity.name)
      |> put_change(:sso_identity_id, identity.id)
      |> put_change(:last_sso_login, NaiveDateTime.utc_now(:second))
      |> put_assoc(:sso_domain, domain)

    with {:ok, user} <-
           Repo.update_with_audit(changeset, "sso_user_provisioned", %{
             team_id: integration.team_id
           }) do
      {:ok, :sso, integration.team, user}
    end
  end

  defp provision_standard_user(user, identity, integration, domain) do
    changeset =
      user
      |> change()
      |> put_change(:type, :sso)
      |> put_change(:name, identity.name)
      |> put_change(:sso_identity_id, identity.id)
      |> put_change(:last_sso_login, NaiveDateTime.utc_now(:second))
      |> put_assoc(:sso_integration, integration)
      |> put_assoc(:sso_domain, domain)

    with :ok <- ensure_team_member(integration.team, user),
         :ok <- ensure_one_membership(user, integration.team),
         :ok <- ensure_empty_personal_team(user, integration.team),
         :ok <- Auth.UserSessions.revoke_all(user),
         {:ok, user} <-
           Repo.update_with_audit(changeset, "sso_user_provisioned", %{
             team_id: integration.team_id
           }) do
      {:ok, :standard, integration.team, user}
    end
  end

  defp provision_identity(identity, integration, domain) do
    random_password =
      64
      |> :crypto.strong_rand_bytes()
      |> Base.encode64(padding: false)

    params = %{
      email: identity.email,
      name: identity.name,
      password: random_password,
      password_confirmation: random_password
    }

    changeset =
      Auth.User.new(params)
      |> put_change(:email_verified, true)
      |> put_change(:type, :sso)
      |> put_change(:sso_identity_id, identity.id)
      |> put_change(:last_sso_login, NaiveDateTime.utc_now(:second))
      |> put_assoc(:sso_integration, integration)
      |> put_assoc(:sso_domain, domain)

    team = integration.team
    role = integration.team.policy.sso_default_role
    now = NaiveDateTime.utc_now(:second)

    result =
      Repo.transaction(fn ->
        with {:ok, user} <-
               Repo.insert_with_audit(changeset, "sso_user_provisioned", %{team_id: team.id}),
             :ok <- Teams.Invitations.check_team_member_limit(team, role, user.email),
             {:ok, team_membership} <-
               Teams.Invitations.create_team_membership(team, role, user, now) do
          if team_membership.role != :guest do
            {:identity, team, user}
          else
            Repo.rollback(:integration_not_found)
          end
        else
          {:error, %{errors: [email: {_, attrs}]}} ->
            true = {:constraint, :unique} in attrs
            Repo.rollback(:integration_not_found)

          {:error, {:over_limit, _}} ->
            Repo.rollback(:over_limit)
        end
      end)

    case result do
      {:ok, {type, team, user}} ->
        {:ok, type, team, user}

      {:error, _} = error ->
        error
    end
  end

  defp ensure_team_member(team, user) do
    case Teams.Memberships.team_role(team, user) do
      {:ok, role} when role != :guest ->
        :ok

      _ ->
        {:error, :integration_not_found}
    end
  end

  defp ensure_one_membership(user, team) do
    if Teams.Users.team_member?(user, except: [team.id], only_setup?: true) do
      {:error, :multiple_memberships, team, user}
    else
      :ok
    end
  end

  defp ensure_empty_personal_team(user, team) do
    case Teams.get_by_owner(user, only_not_setup?: true) do
      {:ok, personal_team} ->
        subscription = Teams.Billing.get_subscription(personal_team)

        no_subscription? =
          is_nil(subscription) or subscription.status == Subscription.Status.deleted()

        zero_sites? = Teams.owned_sites_count(personal_team) == 0

        if no_subscription? and zero_sites? do
          :ok
        else
          {:error, :active_personal_team, team, user}
        end

      {:error, :no_team} ->
        :ok
    end
  end
end