1 files changed, 59 insertions, 29 deletions

diff --git a/app/Controller/Twofactor.php b/app/Controller/Twofactor.php
index aeb13acc..8dbfcf66 100644
--- a/app/Controller/Twofactor.php
+++ b/app/Controller/Twofactor.php
@@ -23,7 +23,7 @@ class Twofactor extends User
     }
 
     /**
-     * Index
+     * Show form to disable/enable 2FA
      *
      * @access public
      */
@@ -31,54 +31,45 @@ class Twofactor extends User
     {
         $user = $this->getUser();
         $this->checkCurrentUser($user);
-
-        $provider = $this->authenticationManager->getPostAuthenticationProvider();
-        $label = $user['email'] ?: $user['username'];
-
-        $provider->setSecret($user['twofactor_secret']);
+        unset($this->sessionStorage->twoFactorSecret);
 
         $this->response->html($this->layout('twofactor/index', array(
             'user' => $user,
-            'qrcode_url' => $user['twofactor_activated'] == 1 ? $provider->getQrCodeUrl($label) : '',
-            'key_url' => $user['twofactor_activated'] == 1 ? $provider->getKeyUrl($label) : '',
+            'provider' => $this->authenticationManager->getPostAuthenticationProvider()->getName(),
         )));
     }
 
     /**
-     * Enable/disable 2FA
+     * Show page with secret and test form
      *
      * @access public
      */
-    public function save()
+    public function show()
     {
         $user = $this->getUser();
         $this->checkCurrentUser($user);
 
-        $values = $this->request->getValues();
+        $label = $user['email'] ?: $user['username'];
+        $provider = $this->authenticationManager->getPostAuthenticationProvider();
 
-        if (isset($values['twofactor_activated']) && $values['twofactor_activated'] == 1) {
-            $this->user->update(array(
-                'id' => $user['id'],
-                'twofactor_activated' => 1,
-                'twofactor_secret' => $this->authenticationManager->getPostAuthenticationProvider()->getSecret(),
-            ));
+        if (! isset($this->sessionStorage->twoFactorSecret)) {
+            $provider->generateSecret();
+            $provider->beforeCode();
+            $this->sessionStorage->twoFactorSecret = $provider->getSecret();
         } else {
-            $this->user->update(array(
-                'id' => $user['id'],
-                'twofactor_activated' => 0,
-                'twofactor_secret' => '',
-            ));
+            $provider->setSecret($this->sessionStorage->twoFactorSecret);
         }
 
-        // Allow the user to test or disable the feature
-        $this->userSession->disablePostAuthentication();
-
-        $this->flash->success(t('User updated successfully.'));
-        $this->response->redirect($this->helper->url->to('twofactor', 'index', array('user_id' => $user['id'])));
+        $this->response->html($this->layout('twofactor/show', array(
+            'user' => $user,
+            'secret' => $this->sessionStorage->twoFactorSecret,
+            'qrcode_url' => $provider->getQrCodeUrl($label),
+            'key_url' => $provider->getKeyUrl($label),
+        )));
     }
 
     /**
-     * Test code
+     * Test code and save secret
      *
      * @access public
      */
@@ -91,14 +82,47 @@ class Twofactor extends User
 
         $provider = $this->authenticationManager->getPostAuthenticationProvider();
         $provider->setCode(empty($values['code']) ? '' : $values['code']);
-        $provider->setSecret($user['twofactor_secret']);
+        $provider->setSecret($this->sessionStorage->twoFactorSecret);
 
         if ($provider->authenticate()) {
             $this->flash->success(t('The two factor authentication code is valid.'));
+
+            $this->user->update(array(
+                'id' => $user['id'],
+                'twofactor_activated' => 1,
+                'twofactor_secret' => $this->authenticationManager->getPostAuthenticationProvider()->getSecret(),
+            ));
+
+            unset($this->sessionStorage->twoFactorSecret);
+            $this->userSession->disablePostAuthentication();
+
+            $this->response->redirect($this->helper->url->to('twofactor', 'index', array('user_id' => $user['id'])));
         } else {
             $this->flash->failure(t('The two factor authentication code is not valid.'));
+            $this->response->redirect($this->helper->url->to('twofactor', 'show', array('user_id' => $user['id'])));
         }
+    }
 
+    /**
+     * Disable 2FA for the current user
+     *
+     * @access public
+     */
+    public function deactivate()
+    {
+        $user = $this->getUser();
+        $this->checkCurrentUser($user);
+
+        $this->user->update(array(
+            'id' => $user['id'],
+            'twofactor_activated' => 0,
+            'twofactor_secret' => '',
+        ));
+
+        // Allow the user to test or disable the feature
+        $this->userSession->disablePostAuthentication();
+
+        $this->flash->success(t('User updated successfully.'));
         $this->response->redirect($this->helper->url->to('twofactor', 'index', array('user_id' => $user['id'])));
     }
 
@@ -135,6 +159,12 @@ class Twofactor extends User
      */
     public function code()
     {
+        if (! isset($this->sessionStorage->twoFactorBeforeCodeCalled)) {
+            $provider = $this->authenticationManager->getPostAuthenticationProvider();
+            $provider->beforeCode();
+            $this->sessionStorage->twoFactorBeforeCodeCalled = true;
+        }
+
         $this->response->html($this->template->layout('twofactor/check', array(
             'title' => t('Check two factor authentication code'),
         )));