Add 2FA to the phx.gen.auth generator

lib/mix/tasks/phx.gen.auth.ex

@@ -131,11 +131,18 @@ defmodule Mix.Tasks.Phx.Gen.Auth do
     {opts, parsed} = OptionParser.parse!(args, strict: @switches)
     validate_args!(parsed)
     hashing_library = build_hashing_library!(opts)
+    ecto_base64_dependency = ~s|{:ecto_base64, "~> 0.1.0"}|
+    totp_dependency = ~s|{:nimble_totp, "~> 1.0"}|
+    qrcode_depencency = ~s|{:eqrcode, "~> 0.1.10"}|
 
     context_args = OptionParser.to_argv(opts, switches: @switches) ++ parsed
     {context, schema} = Gen.Context.build(context_args, __MODULE__)
 
-    context = put_live_option(context)
+    context = 
+      context
+      |> put_live_option()
+      |> put_totp_option()
+
     Gen.Context.prompt_for_code_injection(context)
 
     if "--no-compile" not in args do
@@ -167,7 +174,8 @@ defmodule Mix.Tasks.Phx.Gen.Auth do
       router_scope: router_scope(context),
       web_path_prefix: web_path_prefix(schema),
       test_case_options: test_case_options(ecto_adapter),
-      live?: Keyword.fetch!(context.opts, :live)
+      live?: Keyword.fetch!(context.opts, :live),
+      totp?: Keyword.fetch!(context.opts, :totp)
     ]
 
     paths = Mix.Phoenix.generator_paths()
@@ -179,6 +187,9 @@ defmodule Mix.Tasks.Phx.Gen.Auth do
     |> inject_conn_case_helpers(paths, binding)
     |> inject_config(hashing_library)
     |> maybe_inject_mix_dependency(hashing_library)
+    |> maybe_inject_mix_dependency(ecto_base64_dependency)
+    |> maybe_inject_mix_dependency(totp_dependency)
+    |> maybe_inject_mix_dependency(qrcode_depencency)
     |> inject_routes(paths, binding)
     |> maybe_inject_router_import(binding)
     |> maybe_inject_router_plug()
@@ -257,32 +268,19 @@ defmodule Mix.Tasks.Phx.Gen.Auth do
     web_path = to_string(schema.web_path)
     controller_pre = Path.join([web_pre, "controllers", web_path])
 
-    default_files = [
-      "migration.ex": [migrations_pre, "#{timestamp()}_create_#{schema.table}_auth_tables.exs"],
-      "notifier.ex": [context.dir, "#{singular}_notifier.ex"],
-      "schema.ex": [context.dir, "#{singular}.ex"],
-      "schema_token.ex": [context.dir, "#{singular}_token.ex"],
-      "auth.ex": [web_pre, web_path, "#{singular}_auth.ex"],
-      "auth_test.exs": [web_test_pre, web_path, "#{singular}_auth_test.exs"],
-      "session_controller.ex": [controller_pre, "#{singular}_session_controller.ex"],
-      "session_controller_test.exs": [
-        web_test_pre,
-        "controllers",
-        web_path,
-        "#{singular}_session_controller_test.exs"
-      ]
-    ]
-
-    case Keyword.fetch(context.opts, :live) do
-      {:ok, true} ->
-        live_files = [
-          "registration_live.ex": [
-            web_pre,
-            "live",
-            web_path,
-            "#{singular}_live",
-            "registration.ex"
-          ],
+    live? = Keyword.get(context.opts, :live) == true
+    totp? = Keyword.get(context.opts, :totp) == true
+
+    generated_files = [
+      if totp? do
+        ["totp_controller.ex": [controller_pre, "#{singular}_totp_controller.ex"]]
+      end,
+      if live? do
+        [
+          if totp? do
+            ["totp_live.ex": [web_pre, "live", web_path, "#{singular}_live", "totp.ex"]]
+          end,
+          "registration_live.ex": [web_pre, "live", web_path, "#{singular}_live", "registration.ex"],
           "registration_live_test.exs": [
             web_test_pre,
             "live",
@@ -363,11 +361,12 @@ defmodule Mix.Tasks.Phx.Gen.Auth do
             "confirmation_instructions_test.exs"
           ]
         ]
-
-        remap_files(default_files ++ live_files)
-
-      _ ->
-        non_live_files = [
+      else
+        [
+          if totp? do
+            ["totp_html.ex": [controller_pre, "#{singular}_totp_html.ex"],
+            "totp_new.html.heex": [controller_pre, "#{singular}_totp_html", "new.html.heex"]]
+          end,
           "confirmation_html.ex": [controller_pre, "#{singular}_confirmation_html.ex"],
           "confirmation_new.html.heex": [
             controller_pre,
@@ -422,6 +421,7 @@ defmodule Mix.Tasks.Phx.Gen.Auth do
           ],
           "session_html.ex": [controller_pre, "#{singular}_session_html.ex"],
           "session_new.html.heex": [controller_pre, "#{singular}_session_html", "new.html.heex"],
+
           "settings_html.ex": [web_pre, "controllers", web_path, "#{singular}_settings_html.ex"],
           "settings_controller.ex": [controller_pre, "#{singular}_settings_controller.ex"],
           "settings_edit.html.heex": [
@@ -436,9 +436,26 @@ defmodule Mix.Tasks.Phx.Gen.Auth do
             "#{singular}_settings_controller_test.exs"
           ]
         ]
+      end,
+      "migration.ex": [migrations_pre, "#{timestamp()}_create_#{schema.table}_auth_tables.exs"],
+      "notifier.ex": [context.dir, "#{singular}_notifier.ex"],
+      "schema.ex": [context.dir, "#{singular}.ex"],
+      "schema_token.ex": [context.dir, "#{singular}_token.ex"],
+      "auth.ex": [web_pre, web_path, "#{singular}_auth.ex"],
+      "auth_test.exs": [web_test_pre, web_path, "#{singular}_auth_test.exs"],
+      "session_controller.ex": [controller_pre, "#{singular}_session_controller.ex"],
+      "session_controller_test.exs": [
+        web_test_pre,
+        "controllers",
+        web_path,
+        "#{singular}_session_controller_test.exs"
+      ]
+    ]
 
-        remap_files(default_files ++ non_live_files)
-    end
+    generated_files
+    |> List.flatten()
+    |> Enum.reject(&is_nil/1)
+    |> remap_files()
   end
 
   defp remap_files(files) do
@@ -507,9 +524,11 @@ defmodule Mix.Tasks.Phx.Gen.Auth do
     context
   end
 
-  defp maybe_inject_mix_dependency(%Context{context_app: ctx_app} = context, %HashingLibrary{
-         mix_dependency: mix_dependency
-       }) do
+  defp maybe_inject_mix_dependency(context, %HashingLibrary{} = lib) do
+    maybe_inject_mix_dependency(context, lib.mix_dependency)
+  end
+
+  defp maybe_inject_mix_dependency(%Context{context_app: ctx_app} = context, mix_dependency) do
     file_path = Mix.Phoenix.context_app_path(ctx_app, "mix.exs")
 
     file = File.read!(file_path)
@@ -872,25 +891,30 @@ defmodule Mix.Tasks.Phx.Gen.Auth do
   defp test_case_options(adapter) when is_atom(adapter), do: ""
 
   defp put_live_option(schema) do
-    opts =
-      case Keyword.fetch(schema.opts, :live) do
-        {:ok, _live?} ->
-          schema.opts
+    case Keyword.fetch(schema.opts, :live) do
+      {:ok, _live?} ->
+        schema
 
-        _ ->
-          Mix.shell().info("""
-          An authentication system can be created in two different ways:
-          - Using Phoenix.LiveView (default)
-          - Using Phoenix.Controller only\
-          """)
+      _ ->
+        Mix.shell().info("""
+        An authentication system can be created in two different ways:
+        - Using Phoenix.LiveView (default)
+        - Using Phoenix.Controller only\
+        """)
 
-          if Mix.shell().yes?("Do you want to create a LiveView based authentication system?") do
-            Keyword.put_new(schema.opts, :live, true)
-          else
-            Keyword.put_new(schema.opts, :live, false)
-          end
-      end
+        live? = Mix.shell().yes?("Do you want to create a LiveView based authentication system?")
+        Map.put(schema, :opts, Keyword.put_new(schema.opts, :live, live?))
+    end
+  end
+
+  defp put_totp_option(schema) do
+    case Keyword.fetch(schema.opts, :totp) do
+      {:ok, _totp?} -> 
+        schema
 
-    Map.put(schema, :opts, opts)
+      _ -> 
+        totp? = Mix.shell().yes?("Do you want to generate two-factor authentication too?")
+        Map.put(schema, :opts, Keyword.put_new(schema.opts, :totp, totp?))
+    end
   end
 end

priv/templates/phx.gen.auth/auth.ex

@@ -26,7 +26,9 @@ defmodule <%= inspect auth_module %> do
   if you are not using LiveView.
   """
   def log_in_<%= schema.singular %>(conn, <%= schema.singular %>, params \\ %{}) do
-    token = <%= inspect context.alias %>.generate_<%= schema.singular %>_session_token(<%= schema.singular %>)
+    <%= if totp? do %><%= inspect context.alias %>.mark_<%= schema.singular %>_login(<%= schema.singular %>)
+
+    <% end %>token = <%= inspect context.alias %>.generate_<%= schema.singular %>_session_token(<%= schema.singular %>)
     <%= schema.singular %>_return_to = get_session(conn, :<%= schema.singular %>_return_to)
 
     conn

priv/templates/phx.gen.auth/context_functions.ex

@@ -341,4 +341,78 @@
       {:ok, %{<%= schema.singular %>: <%= schema.singular %>}} -> {:ok, <%= schema.singular %>}
       {:error, :<%= schema.singular %>, changeset, _} -> {:error, changeset}
     end
+  end<%= if totp? do %>
+
+  @doc """
+  Checks if the entered OTP code is valid or not.
+
+  We're also allowing codes to be 30 seconds in the past or future,
+  to account for slightly mismatching times on different devices.
+  """
+  def valid_<%= schema.singular %>_totp?(<%= schema.singular %>_or_secret, validation_code, offset \\ 30, opts \\ [])
+
+  def valid_<%= schema.singular %>_totp?(%<%= inspect schema.alias %>{} = <%= schema.singular %>, validation_code, offset, opts)
+      when is_binary(validation_code) and is_binary(<%= schema.singular %>.totp_secret) do
+    opts = Keyword.put_new(opts, :since, <%= schema.singular %>.last_login)
+    valid_<%= schema.singular %>_totp?(<%= schema.singular %>.totp_secret, validation_code, offset, opts)
+  end
+
+  def valid_<%= schema.singular %>_totp?(totp_secret, validation_code, offset, opts) when is_binary(validation_code) do
+    Enum.any?([-offset, 0, offset], fn offset ->
+      time = Keyword.get(opts, :time, System.os_time(:second))
+      opts = Keyword.put(opts, :time, time + offset)
+
+      NimbleTOTP.valid?(totp_secret, validation_code, opts)
+    end)
+  end
+
+  @doc """
+  Returns an `%Ecto.Changeset{}` for changing the <%= schema.singular %> OTP secret.
+
+  ## Examples
+
+      iex> change_<%= schema.singular %>_totp(<%= schema.singular %>)
+      %Ecto.Changeset{data: %<%= inspect schema.alias %>{}}
+
+  """
+  def change_<%= schema.singular %>_totp(<%= schema.singular %>, attrs \\ %{}) do
+    <%= inspect schema.alias %>.totp_changeset(<%= schema.singular %>, attrs)
   end
+
+  @doc """
+  Updates the `last_login` field for the given <%= schema.singular %>.
+  """
+  def mark_<%= schema.singular %>_login(<%= schema.singular %>) do
+    <%= schema.singular %>
+    |> <%= inspect schema.alias %>.login_changeset()
+    |> Repo.update!()
+  end
+
+  @doc """
+  Enables 2FA for an account if the provided code is valid for
+  the given secret.
+  """
+  def enable_<%= schema.singular %>_totp(<%= schema.singular %>, secret, code) do
+    attrs = %{totp_secret: secret}
+
+    <%= schema.singular %>
+    |> <%= inspect schema.alias %>.totp_changeset(attrs)
+    |> <%= inspect schema.alias %>.validate_totp(code)
+    |> <%= inspect schema.alias %>.login_changeset()
+    |> Repo.update()
+  end
+
+  @doc """
+  Disables 2FA for an account if the provided code is valid for
+  the OTP secret on the account.
+  """
+  def disable_<%= schema.singular %>_totp(<%= schema.singular %>, password, code) do
+    attrs = %{totp_secret: nil}
+
+    <%= schema.singular %>
+    |> <%= inspect schema.alias %>.totp_changeset(attrs)
+    |> <%= inspect schema.alias %>.validate_totp(code, for: :<%= schema.singular %>)
+    |> <%= inspect schema.alias %>.validate_current_password(password)
+    |> <%= inspect schema.alias %>.login_changeset()
+    |> Repo.update()
+  end<% end %>

priv/templates/phx.gen.auth/login_live.ex

@@ -43,7 +43,9 @@ defmodule <%= inspect context.web_module %>.<%= inspect Module.concat(schema.web
 
   def mount(_params, _session, socket) do
     email = Phoenix.Flash.get(socket.assigns.flash, :email)
-    form = to_form(%{"email" => email}, as: "<%= schema.singular %>")
-    {:ok, assign(socket, form: form), temporary_assigns: [form: form]}
+    remember_me = Phoenix.Flash.get(socket.assigns.flash, :remember_me)
+    form = to_form(%{"email" => email, "remember_me" => remember_me}, as: "<%= schema.singular %>")
+
+    {:ok, assign(socket, :form, form), temporary_assigns: [form: form]}
   end
 end

priv/templates/phx.gen.auth/migration.ex

@@ -8,6 +8,8 @@ defmodule <%= inspect schema.repo %>.Migrations.Create<%= Macro.camelize(schema.
 <%= if schema.binary_id do %>      add :id, :binary_id, primary_key: true
 <% end %>      <%= migration.column_definitions[:email] %>
       add :hashed_password, :string, null: false
+      <%= if totp? do %>add :totp_secret, :string
+  	  <% end %>add :last_login, :naive_datetime
       add :confirmed_at, <%= inspect schema.timestamp_type %>
 
       timestamps(<%= if schema.timestamp_type != :naive_datetime, do: "type: #{inspect schema.timestamp_type}" %>)

priv/templates/phx.gen.auth/routes.ex

@@ -8,17 +8,21 @@
       on_mount: [{<%= inspect auth_module %>, :redirect_if_<%= schema.singular %>_is_authenticated}] do
       live "/<%= schema.plural %>/register", <%= inspect schema.alias %>Live.Registration, :new
       live "/<%= schema.plural %>/log-in", <%= inspect schema.alias %>Live.Login, :new
-      live "/<%= schema.plural %>/reset-password", <%= inspect schema.alias %>Live.ForgotPassword, :new
+      <%= if totp? do %>live "/<%= schema.plural %>/2fa", <%= inspect schema.alias %>Live.TOTP, :new
+      <% end %>live "/<%= schema.plural %>/reset-password", <%= inspect schema.alias %>Live.ForgotPassword, :new
       live "/<%= schema.plural %>/reset-password/:token", <%= inspect schema.alias %>Live.ResetPassword, :edit
     end
 
-    post "/<%= schema.plural %>/log-in", <%= inspect schema.alias %>SessionController, :create<% else %>
+    post "/<%= schema.plural %>/log-in", <%= inspect schema.alias %>SessionController, :create<%= if totp? do %>
+    post "/<%= schema.plural %>/2fa", <%= inspect schema.alias %>TOTPController, :create<% end %><% else %>
 
     get "/<%= schema.plural %>/register", <%= inspect schema.alias %>RegistrationController, :new
     post "/<%= schema.plural %>/register", <%= inspect schema.alias %>RegistrationController, :create
     get "/<%= schema.plural %>/log-in", <%= inspect schema.alias %>SessionController, :new
     post "/<%= schema.plural %>/log-in", <%= inspect schema.alias %>SessionController, :create
-    get "/<%= schema.plural %>/reset-password", <%= inspect schema.alias %>ResetPasswordController, :new
+    <%= if totp? do %>get "/<%= schema.plural %>/2fa", <%= inspect schema.alias %>TOTPController, :new
+    post "/<%= schema.plural %>/2fa", <%= inspect schema.alias %>TOTPController, :create
+    <% end %>get "/<%= schema.plural %>/reset-password", <%= inspect schema.alias %>ResetPasswordController, :new
     post "/<%= schema.plural %>/reset-password", <%= inspect schema.alias %>ResetPasswordController, :create
     get "/<%= schema.plural %>/reset-password/:token", <%= inspect schema.alias %>ResetPasswordController, :edit
     put "/<%= schema.plural %>/reset-password/:token", <%= inspect schema.alias %>ResetPasswordController, :update<% end %>