From 742fb0d27b1a24a67b51cf7d4cc85ca9d5328bfb Mon Sep 17 00:00:00 2001
From: aignatov-bio <aignatov@biosistemika.com>
Date: Wed, 1 Jul 2020 14:41:55 +0200
Subject: [PATCH] Small 2fa improvments

---
 app/controllers/users/registrations_controller.rb | 7 +++----
 app/controllers/users/sessions_controller.rb      | 2 +-
 app/models/user.rb                                | 8 ++++++++
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/app/controllers/users/registrations_controller.rb b/app/controllers/users/registrations_controller.rb
index 67801ccaf..1de739f67 100644
--- a/app/controllers/users/registrations_controller.rb
+++ b/app/controllers/users/registrations_controller.rb
@@ -183,9 +183,8 @@ class Users::RegistrationsController < Devise::RegistrationsController
   end
 
   def two_factor_enable
-    totp = ROTP::TOTP.new(current_user.otp_secret, issuer: 'SciNote')
-    if totp.verify(params[:submit_code], drift_behind: 10)
-      current_user.update!(two_factor_auth_enabled: true)
+    if current_user.valid_otp?(params[:submit_code])
+      current_user.enable_2fa
       redirect_to edit_user_registration_path
     else
       render json: { error: t('users.registrations.edit.2fa_errors.wrong_submit_code') }, status: :unprocessable_entity
@@ -194,7 +193,7 @@ class Users::RegistrationsController < Devise::RegistrationsController
 
   def two_factor_disable
     if current_user.valid_password?(params[:password])
-      current_user.update!(two_factor_auth_enabled: false, otp_secret: nil)
+      current_user.disable_2fa
       redirect_to edit_user_registration_path
     else
       render json: { error: t('users.registrations.edit.2fa_errors.wrong_password') }, status: :forbidden
diff --git a/app/controllers/users/sessions_controller.rb b/app/controllers/users/sessions_controller.rb
index 50d7c7258..aef23a8c1 100644
--- a/app/controllers/users/sessions_controller.rb
+++ b/app/controllers/users/sessions_controller.rb
@@ -4,7 +4,7 @@ class Users::SessionsController < Devise::SessionsController
   layout :session_layout
 
   # before_filter :configure_sign_in_params, only: [:create]
-  after_action :after_sign_in, only: :create
+  after_action :after_sign_in, only: %i(create authenticate_with_two_factor)
   prepend_before_action :redirect_2fa, only: :create
 
   rescue_from ActionController::InvalidAuthenticityToken do
diff --git a/app/models/user.rb b/app/models/user.rb
index eed3623c0..c8e400c23 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -635,6 +635,14 @@ class User < ApplicationRecord
     save!
   end
 
+  def enable_2fa
+    update!(two_factor_auth_enabled: true)
+  end
+
+  def disable_2fa
+    update!(two_factor_auth_enabled: false, otp_secret: nil)
+  end
+
   protected
 
   def confirmation_required?