From 9833f6104381a2318eaf23e0e48d96196527a425 Mon Sep 17 00:00:00 2001
From: Graham McIntire <graham@mcintire.me>
Date: Wed, 28 Jan 2026 15:18:08 -0600
Subject: [PATCH] login redirect fix

---
 lib/towerops_web/components/layouts.ex        |  8 -----
 .../components/marketing_layouts.ex           |  4 ---
 .../live/account_live/totp_enrollment.ex      | 33 +++++++++++++------
 priv/static/changelog.txt                     |  1 +
 .../account_live/totp_enrollment_test.exs     | 27 ++++++++++++++-
 5 files changed, 50 insertions(+), 23 deletions(-)

diff --git a/lib/towerops_web/components/layouts.ex b/lib/towerops_web/components/layouts.ex
index 2b566744..476f03c9 100644
--- a/lib/towerops_web/components/layouts.ex
+++ b/lib/towerops_web/components/layouts.ex
@@ -48,10 +48,6 @@ defmodule ToweropsWeb.Layouts do
     requires_cookie_consent = Process.get(:requires_cookie_consent, false)
     assigns = Map.put(assigns, :requires_cookie_consent, requires_cookie_consent)
 
-    Logger.info(
-      "Layout app/1: requires_cookie_consent = #{inspect(assigns.requires_cookie_consent)}, process dict = #{inspect(Process.get(:requires_cookie_consent))}"
-    )
-
     ~H"""
     <div class="flex min-h-screen flex-col bg-gray-50 dark:bg-gray-950">
       <header class="border-b border-gray-200 bg-white dark:border-white/10 dark:bg-gray-900">
@@ -614,10 +610,6 @@ defmodule ToweropsWeb.Layouts do
     requires_cookie_consent = Process.get(:requires_cookie_consent, false)
     assigns = Map.put(assigns, :requires_cookie_consent, requires_cookie_consent)
 
-    Logger.info(
-      "Layout admin/1: requires_cookie_consent = #{inspect(assigns.requires_cookie_consent)}, process dict = #{inspect(Process.get(:requires_cookie_consent))}"
-    )
-
     ~H"""
     <div class="flex min-h-screen flex-col bg-gray-50 dark:bg-gray-950">
       <nav class="bg-white border-b border-gray-200 dark:border-white/10 dark:bg-gray-900">
diff --git a/lib/towerops_web/components/marketing_layouts.ex b/lib/towerops_web/components/marketing_layouts.ex
index 47359ea4..f00cbabf 100644
--- a/lib/towerops_web/components/marketing_layouts.ex
+++ b/lib/towerops_web/components/marketing_layouts.ex
@@ -29,10 +29,6 @@ defmodule ToweropsWeb.MarketingLayouts do
     requires_cookie_consent = Process.get(:requires_cookie_consent, false)
     assigns = Map.put(assigns, :requires_cookie_consent, requires_cookie_consent)
 
-    Logger.info(
-      "Layout marketing/1: requires_cookie_consent = #{inspect(assigns.requires_cookie_consent)}, process dict = #{inspect(Process.get(:requires_cookie_consent))}"
-    )
-
     ~H"""
     <div class="bg-white">
       <header class="py-10">
diff --git a/lib/towerops_web/live/account_live/totp_enrollment.ex b/lib/towerops_web/live/account_live/totp_enrollment.ex
index 1eb1138a..d316c354 100644
--- a/lib/towerops_web/live/account_live/totp_enrollment.ex
+++ b/lib/towerops_web/live/account_live/totp_enrollment.ex
@@ -37,21 +37,34 @@ defmodule ToweropsWeb.AccountLive.TotpEnrollment do
     secret = socket.assigns.secret
 
     if Accounts.verify_totp(secret, code) do
-      case Accounts.enable_totp(user, secret) do
-        {:ok, _user} ->
-          {:noreply,
-           socket
-           |> put_flash(:info, "Two-factor authentication enabled successfully!")
-           |> redirect(to: ~p"/devices")}
-
-        {:error, _changeset} ->
-          {:noreply, assign(socket, :error, "Failed to enable two-factor authentication. Please try again.")}
-      end
+      handle_valid_code(socket, user, secret)
     else
       {:noreply, assign(socket, :error, "Invalid code. Please try again.")}
     end
   end
 
+  defp handle_valid_code(socket, user, secret) do
+    case Accounts.enable_totp(user, secret) do
+      {:ok, updated_user} ->
+        redirect_path = get_post_enrollment_path(updated_user)
+
+        {:noreply,
+         socket
+         |> put_flash(:info, "Two-factor authentication enabled successfully!")
+         |> redirect(to: redirect_path)}
+
+      {:error, _changeset} ->
+        {:noreply, assign(socket, :error, "Failed to enable two-factor authentication. Please try again.")}
+    end
+  end
+
+  defp get_post_enrollment_path(user) do
+    case Towerops.Organizations.list_user_organizations(user.id) do
+      [_first_org | _] -> ~p"/devices"
+      [] -> ~p"/orgs"
+    end
+  end
+
   @impl true
   def render(assigns) do
     ~H"""
diff --git a/priv/static/changelog.txt b/priv/static/changelog.txt
index 27ee1ee7..3a23a6f0 100644
--- a/priv/static/changelog.txt
+++ b/priv/static/changelog.txt
@@ -9,6 +9,7 @@ Devices Working
 * GDPR: API documentation for account data export
 * GDPR: My Data links in user settings and navigation
 * Bug fix: SNMP Community when using remote agent
+* Feature: Add mandatory TOTP MFA
 
 2025-01-27
 * Infrastructure improvements
diff --git a/test/towerops_web/live/account_live/totp_enrollment_test.exs b/test/towerops_web/live/account_live/totp_enrollment_test.exs
index 217d8599..9a325e1d 100644
--- a/test/towerops_web/live/account_live/totp_enrollment_test.exs
+++ b/test/towerops_web/live/account_live/totp_enrollment_test.exs
@@ -25,6 +25,9 @@ defmodule ToweropsWeb.AccountLive.TotpEnrollmentTest do
     end
 
     test "accepts valid TOTP code and enables TOTP", %{conn: conn, user: user} do
+      # Create an organization for the user so they redirect to /devices
+      {:ok, _org} = Towerops.Organizations.create_organization(%{name: "Test Org"}, user.id)
+
       {:ok, view, _html} = live(conn, ~p"/account/totp-enrollment")
 
       # Get the secret from the view's assigns (via testing API)
@@ -38,7 +41,7 @@ defmodule ToweropsWeb.AccountLive.TotpEnrollmentTest do
       |> form("form", %{code: code})
       |> render_submit()
 
-      # Should redirect to devices page
+      # Should redirect to devices page (user has organization)
       assert_redirect(view, ~p"/devices")
 
       # Verify TOTP was enabled in database
@@ -47,6 +50,28 @@ defmodule ToweropsWeb.AccountLive.TotpEnrollmentTest do
       assert updated_user.totp_secret == secret
     end
 
+    test "redirects to orgs page when user has no organizations", %{conn: conn, user: user} do
+      {:ok, view, _html} = live(conn, ~p"/account/totp-enrollment")
+
+      # Get the secret from the view's assigns
+      secret = :sys.get_state(view.pid).socket.assigns.secret
+
+      # Generate a valid code
+      code = NimbleTOTP.verification_code(secret)
+
+      # Submit the form
+      view
+      |> form("form", %{code: code})
+      |> render_submit()
+
+      # Should redirect to orgs page (user has no organizations)
+      assert_redirect(view, ~p"/orgs")
+
+      # Verify TOTP was enabled in database
+      updated_user = Accounts.get_user!(user.id)
+      assert Accounts.totp_enabled?(updated_user)
+    end
+
     test "rejects invalid TOTP code", %{conn: conn} do
       {:ok, view, _html} = live(conn, ~p"/account/totp-enrollment")