diff --git a/lib/two_factor_authentication/models/two_factor_authenticatable.rb b/lib/two_factor_authentication/models/two_factor_authenticatable.rb index bb7a8e0..5421369 100644 --- a/lib/two_factor_authentication/models/two_factor_authenticatable.rb +++ b/lib/two_factor_authentication/models/two_factor_authenticatable.rb @@ -33,19 +33,20 @@ module Devise end def authenticate_totp(code, options = {}) - raise "authenticate_totp called with no otp_secret_key set" if otp_secret_key.nil? - totp = ROTP::TOTP.new( - otp_secret_key, digits: options[:otp_length] || self.class.otp_length - ) + 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 - + raise "authenticate_totp called with no otp_secret_key set" if totp_secret.nil? + totp = ROTP::TOTP.new(totp_secret, digits: digits) totp.verify_with_drift(code, drift) end def provisioning_uri(account = nil, options = {}) - raise "provisioning_uri called with no otp_secret_key set" if otp_secret_key.nil? + totp_secret = options[:otp_secret_key] || otp_secret_key + options[:digits] ||= options[: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(otp_secret_key, options).provisioning_uri(account) + ROTP::TOTP.new(totp_secret, options).provisioning_uri(account) end def need_two_factor_authentication?(request) @@ -73,8 +74,14 @@ module Devise respond_to?(:otp_secret_key) && !otp_secret_key.nil? end + def confirm_totp_secret(secret, code, options = {}) + return false unless authenticate_totp(code, {otp_secret_key: secret}) + self.otp_secret_key = secret + true + end + def generate_totp_secret - self.otp_secret_key = ROTP::Base32.random_base32 + ROTP::Base32.random_base32 end def create_direct_otp(options = {}) diff --git a/spec/lib/two_factor_authentication/models/two_factor_authenticatable_spec.rb b/spec/lib/two_factor_authentication/models/two_factor_authenticatable_spec.rb index 8d69a7a..b3af895 100644 --- a/spec/lib/two_factor_authentication/models/two_factor_authenticatable_spec.rb +++ b/spec/lib/two_factor_authentication/models/two_factor_authenticatable_spec.rb @@ -122,7 +122,7 @@ describe Devise::Models::TwoFactorAuthenticatable do describe 'with secret set' do before do instance.email = 'houdini@example.com' - instance.generate_totp_secret + instance.otp_secret_key = instance.generate_totp_secret end it "returns uri with user's email" do @@ -157,13 +157,10 @@ describe Devise::Models::TwoFactorAuthenticatable do shared_examples 'generate_totp_secret' do |klass| let(:instance) { klass.new } - it 'populates otp_secret_key column' do - original_key = instance.otp_secret_key + it 'returns a 16 character string' do + secret = instance.generate_totp_secret - instance.generate_totp_secret - - expect(instance.otp_secret_key).to match(/\w{16}/) - expect(instance.otp_secret_key).to_not eq(original_key) + expect(secret).to match(/\w{16}/) end end @@ -171,6 +168,37 @@ describe Devise::Models::TwoFactorAuthenticatable do it_behaves_like 'generate_totp_secret', EncryptedUser end + describe '#confirm_totp_secret' do + shared_examples 'confirm_totp_secret' do |klass| + let(:instance) { klass.new } + let(:secret) { instance.generate_totp_secret } + let(:totp_helper) { TotpHelper.new(secret, instance.class.otp_length) } + + it 'populates otp_secret_key column when given correct code' do + instance.confirm_totp_secret(secret, totp_helper.totp_code) + + expect(instance.otp_secret_key).to match(secret) + end + + it 'does not populate otp_secret_key when when given incorrect code' do + instance.confirm_totp_secret(secret, '123') + expect(instance.otp_secret_key).to be_nil + end + + it 'returns true when given correct code' do + expect(instance.confirm_totp_secret(secret, totp_helper.totp_code)).to be true + end + + it 'returns false when given incorrect code' do + expect(instance.confirm_totp_secret(secret, '123')).to be false + end + + end + + it_behaves_like 'confirm_totp_secret', GuestUser + it_behaves_like 'confirm_totp_secret', EncryptedUser + end + describe '#max_login_attempts' do let(:instance) { build_guest_user }