From 06c67df575d99fd51a62427d696d6e1211d4b7b5 Mon Sep 17 00:00:00 2001
From: Sam Clegg <sclegg@omb.eop.gov>
Date: Mon, 13 Jun 2016 14:55:49 -0400
Subject: [PATCH] Seperate totp secret generation from confirmation

For most use cases the totp secret needs to be transmitted
to the end user so it its helpful to be able to generate
it, before confirming it.
---
 .../models/two_factor_authenticatable.rb      | 23 ++++++----
 .../models/two_factor_authenticatable_spec.rb | 42 +++++++++++++++----
 2 files changed, 50 insertions(+), 15 deletions(-)

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 }