towerops/lib/towerops_web/controllers/user_session_controller.ex

defmodule ToweropsWeb.UserSessionController do
  @moduledoc false

  use ToweropsWeb, :controller

  alias Towerops.Accounts
  alias ToweropsWeb.UserAuth

  def new(conn, params) do
    # Check for email in query params (from registration redirect) or current user
    email = params["email"] || get_in(conn.assigns, [:current_scope, Access.key(:user), Access.key(:email)])
    form = Phoenix.Component.to_form(%{"email" => email}, as: "user")

    render(conn, :new, form: form)
  end

  # magic link login
  def create(conn, %{"user" => %{"token" => token} = user_params} = params) do
    case Accounts.login_user_by_magic_link(token) do
      {:ok, {user, _expired_tokens}} ->
        conn =
          if params["_action"] == "confirmed" do
            put_flash(conn, :info, t("User confirmed successfully."))
          else
            conn
          end

        UserAuth.log_in_user(conn, user, user_params)

      {:error, :not_found} ->
        conn
        |> put_flash(:error, t("The link is invalid or it has expired."))
        |> render(:new, form: Phoenix.Component.to_form(%{}, as: "user"))
    end
  end

  # email + password login
  def create(conn, %{"user" => %{"email" => email, "password" => password} = user_params}) do
    if user = Accounts.get_user_by_email_and_password(email, password) do
      cond do
        is_nil(user.confirmed_at) ->
          conn
          |> put_flash(
            :error,
            "Please verify your email address before logging in. Check your inbox for a confirmation link."
          )
          |> render(:new, form: Phoenix.Component.to_form(user_params, as: "user"))

        # Check if user has TOTP enabled - require verification before creating session
        Accounts.totp_enabled?(user) ->
          remember_me = Map.get(user_params, "remember_me") == "true"

          conn
          |> put_session(:pending_totp_user_id, user.id)
          |> put_session(:pending_totp_remember_me, remember_me)
          |> redirect(to: ~p"/users/log-in/totp")

        true ->
          # No TOTP, proceed with normal login
          UserAuth.log_in_user(conn, user, user_params)
      end
    else
      form = Phoenix.Component.to_form(user_params, as: "user")

      # In order to prevent user enumeration attacks, don't disclose whether the email is registered.
      conn
      |> put_flash(:error, t("Invalid email or password"))
      |> render(:new, form: form)
    end
  end

  # magic link request
  def create(conn, %{"user" => %{"email" => email}}) do
    _ =
      if user = Accounts.get_user_by_email(email) do
        Accounts.deliver_login_instructions(
          user,
          &url(~p"/users/log-in/#{&1}")
        )
      end

    info =
      "If your email is in our system, you will receive instructions for logging in shortly."

    conn
    |> put_flash(:info, info)
    |> redirect(to: ~p"/users/log-in")
  end

  def confirm(conn, %{"token" => token}) do
    # Check if user was previously confirmed before logging in
    was_confirmed =
      case Accounts.get_user_by_magic_link_token(token) do
        %Accounts.User{confirmed_at: confirmed_at} when not is_nil(confirmed_at) -> true
        _ -> false
      end

    case Accounts.login_user_by_magic_link(token) do
      {:ok, {user, _expired_tokens}} ->
        conn =
          if was_confirmed do
            conn
          else
            put_flash(conn, :info, t("User confirmed successfully."))
          end

        # Always use remember_me: true to keep user logged in
        UserAuth.log_in_user(conn, user, %{"remember_me" => "true"})

      {:error, :not_found} ->
        conn
        |> put_flash(:error, t("Magic link is invalid or it has expired."))
        |> redirect(to: ~p"/users/log-in")
    end
  end

  def switch_org(conn, %{"org_id" => org_id}) do
    user = conn.assigns.current_scope.user
    membership = Towerops.Organizations.get_membership(org_id, user.id)

    if membership do
      {:ok, _user} = Accounts.update_user_profile(user, %{default_organization_id: org_id})

      conn
      |> put_session(:current_organization_id, org_id)
      |> redirect(to: ~p"/dashboard")
    else
      conn
      |> put_flash(:error, t("You don't have access to this organization."))
      |> redirect(to: ~p"/orgs")
    end
  end

  def delete(conn, _params) do
    conn
    |> put_flash(:info, t("Logged out successfully."))
    |> UserAuth.log_out_user()
  end

  def totp_verification_form(conn, _params) do
    case get_session(conn, :pending_totp_user_id) do
      nil ->
        conn
        |> put_flash(:error, t("Please log in first."))
        |> redirect(to: ~p"/users/log-in")

      _user_id ->
        form = Phoenix.Component.to_form(%{}, as: "user")
        render(conn, :totp, form: form)
    end
  end

  def verify_totp(conn, %{"user" => %{"totp_code" => code}}) do
    case get_session(conn, :pending_totp_user_id) do
      nil ->
        redirect_to_login(conn)

      user_id ->
        handle_totp_verification(conn, user_id, code)
    end
  end

  defp redirect_to_login(conn) do
    conn
    |> put_flash(:error, t("Please log in first."))
    |> redirect(to: ~p"/users/log-in")
  end

  defp handle_totp_verification(conn, user_id, code) do
    user = Accounts.get_user!(user_id)

    case Accounts.verify_user_mfa(user, code) do
      {:ok, user, :totp} ->
        complete_totp_login(conn, user)

      {:ok, user, :recovery_code} ->
        conn
        |> put_flash(
          :warning,
          "You used a recovery code. Consider regenerating codes from your account settings."
        )
        |> complete_totp_login(user)

      {:error, _reason} ->
        render_totp_error(conn, code)
    end
  end

  defp complete_totp_login(conn, user) do
    remember_me = get_session(conn, :pending_totp_remember_me)

    user_params =
      if remember_me do
        %{"remember_me" => "true"}
      else
        %{}
      end

    conn
    |> delete_session(:pending_totp_user_id)
    |> delete_session(:pending_totp_remember_me)
    |> UserAuth.log_in_user(user, user_params)
  end

  defp render_totp_error(conn, code) do
    form = Phoenix.Component.to_form(%{"totp_code" => code}, as: "user")

    conn
    |> put_flash(:error, t("Invalid authentication code. Please try again."))
    |> render(:totp, form: form)
  end
end