Improvements to configuration

README.md

@@ -95,7 +95,7 @@ config.otp_length = 6  # TOTP code length
 config.direct_otp_valid_for = 5.minutes  # Time before direct OTP becomes invalid
 config.direct_otp_length = 6  # Direct OTP code length
 config.remember_otp_session_for_seconds = 30.days  # Time before browser has to perform 2fA again. Default is 0.
-config.otp_secret_encryption_key = ENV['OTP_SECRET_ENCRYPTION_KEY']
+config.otp_secret_encryption_key = ENV['OTP_SECRET_ENCRYPTION_KEY'] # Can be also a lambda
 config.second_factor_resource_id = 'id' # Field or method name used to set value for 2fA remember cookie
 config.delete_cookie_on_logout = false # Delete cookie when user signs out, to force 2fA again on login
 ```

app/controllers/devise/two_factor_authentication_controller.rb

@@ -42,7 +42,11 @@ def after_two_factor_success_for(resource)
   end
 
   def set_remember_two_factor_cookie(resource)
-    expires_seconds = resource.class.remember_otp_session_for_seconds
+    expires_seconds = if resource.respond_to?(:remember_otp_session_for_seconds)
+      resource.remember_otp_session_for_seconds
+    else
+      resource.class.remember_otp_session_for_seconds
+    end
 
     if expires_seconds && expires_seconds > 0
       cookies.signed[TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME] = {

lib/two_factor_authentication.rb

@@ -24,14 +24,23 @@ module Devise
   mattr_accessor :remember_otp_session_for_seconds
   @@remember_otp_session_for_seconds = 0
 
-  mattr_accessor :otp_secret_encryption_key
-  @@otp_secret_encryption_key = ''
-
   mattr_accessor :second_factor_resource_id
   @@second_factor_resource_id = 'id'
 
   mattr_accessor :delete_cookie_on_logout
   @@delete_cookie_on_logout = false
+
+  mattr_writer :otp_secret_encryption_key
+  @@otp_secret_encryption_key = ''
+
+  def self.otp_secret_encryption_key
+    if @@otp_secret_encryption_key.respond_to?(:call)
+      @@otp_secret_encryption_key.call
+    else
+      @@otp_secret_encryption_key
+    end
+  end
+  delegate :otp_secret_encryption_key, to: 'self.class'
 end
 
 module TwoFactorAuthentication

lib/two_factor_authentication/hooks/two_factor_authenticatable.rb

@@ -1,7 +1,8 @@
 Warden::Manager.after_authentication do |user, auth, options|
-  if auth.env["action_dispatch.cookies"]
+  cookie_jar = auth.cookies || auth.env["action_dispatch.cookies"]
+  if cookie_jar
     expected_cookie_value = "#{user.class}-#{user.public_send(Devise.second_factor_resource_id)}"
-    actual_cookie_value = auth.env["action_dispatch.cookies"].signed[TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME]
+    actual_cookie_value = cookie_jar.signed[TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME]
     bypass_by_cookie = actual_cookie_value == expected_cookie_value
   end
 
@@ -13,5 +14,11 @@
 end
 
 Warden::Manager.before_logout do |user, auth, _options|
-  auth.cookies.delete TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME if Devise.delete_cookie_on_logout
+  should_delete = Devise.delete_cookie_on_logout
+
+  if user.respond_to?(:delete_cookie_on_logout?)
+    should_delete = user.delete_cookie_on_logout
+  end
+
+  auth.cookies.delete TwoFactorAuthentication::REMEMBER_TFA_COOKIE_NAME if should_delete
 end

lib/two_factor_authentication/models/two_factor_authenticatable.rb

@@ -35,8 +35,8 @@ def authenticate_direct_otp(code)
 
         def authenticate_totp(code, options = {})
           totp_secret = options[:otp_secret_key] || otp_secret_key
-          digits = options[:otp_length] || self.class.otp_length
-          drift = options[:drift] || self.class.allowed_otp_drift_seconds
+          digits = options[:otp_length] || (self.respond_to?(:otp_length) && self.otp_length) || self.class.otp_length
+          drift = options[:drift] || (self.respond_to?(:allowed_otp_drift_seconds) && self.allowed_otp_drift_seconds) || self.class.allowed_otp_drift_seconds
           raise "authenticate_totp called with no otp_secret_key set" if totp_secret.nil?
           totp = ROTP::TOTP.new(totp_secret, digits: digits)
           new_timestamp = totp.verify(
@@ -50,7 +50,7 @@ def authenticate_totp(code, options = {})
 
         def provisioning_uri(account = nil, options = {})
           totp_secret = options[:otp_secret_key] || otp_secret_key
-          options[:digits] ||= options[:otp_length] || self.class.otp_length
+          options[:digits] ||= options[:otp_length] || (self.respond_to?(:otp_length) && self.otp_length) || self.class.otp_length
           raise "provisioning_uri called with no otp_secret_key set" if totp_secret.nil?
           account ||= email if respond_to?(:email)
           ROTP::TOTP.new(totp_secret, options).provisioning_uri(account)
@@ -74,11 +74,15 @@ def send_two_factor_authentication_code(code)
         end
 
         def max_login_attempts?
-          second_factor_attempts_count.to_i >= max_login_attempts.to_i
+          second_factor_attempts_count.to_i > max_login_attempts.to_i
         end
 
         def max_login_attempts
-          self.class.max_login_attempts
+          self.max_login_attempts
+        end
+
+        def attempts_left
+          max_login_attempts.to_i - second_factor_attempts_count.to_i
         end
 
         def totp_enabled?
@@ -100,7 +104,7 @@ def generate_totp_secret
 
         def create_direct_otp(options = {})
           # Create a new random OTP and store it in the database
-          digits = options[:length] || self.class.direct_otp_length || 6
+          digits = options[:length] || (self.respond_to?(:direct_otp_length) && self.direct_otp_length) || self.class.direct_otp_length || 6
           update_attributes(
             direct_otp: random_base10(digits),
             direct_otp_sent_at: Time.now.utc
@@ -118,7 +122,7 @@ def random_base10(digits)
         end
 
         def direct_otp_expired?
-          Time.now.utc > direct_otp_sent_at + self.class.direct_otp_valid_for
+          Time.now.utc > direct_otp_sent_at + self.direct_otp_valid_for
         end
 
         def clear_direct_otp
@@ -166,13 +170,21 @@ def otp_encrypt(value)
         def encryption_options_for(value)
           {
             value: value,
-            key: Devise.otp_secret_encryption_key,
+            key: otp_secret_encryption_key,
             iv: iv_for_attribute,
             salt: salt_for_attribute,
             algorithm: 'aes-256-cbc'
           }
         end
 
+        def otp_secret_encryption_key
+          if self.respond_to?(:otp_secret_encryption_key)
+            self.otp_secret_encryption_key
+          else
+            Devise.otp_secret_encryption_key
+          end
+        end
+
         def iv_for_attribute(algorithm = 'aes-256-cbc')
           iv = encrypted_otp_secret_key_iv
 

lib/two_factor_authentication/version.rb

@@ -1,3 +1,3 @@
 module TwoFactorAuthentication
-  VERSION = "2.2.0".freeze
+  VERSION = "2.3.0".freeze
 end