Improve 2FA

2 files changed, 5 insertions, 3 deletions

diff --git a/app/Controller/Base.php b/app/Controller/Base.php
index f498c3ce..b40e69c0 100644
--- a/app/Controller/Base.php
+++ b/app/Controller/Base.php
@@ -207,10 +207,9 @@ abstract class Base
      */
     public function handle2FA($controller, $action)
     {
-        $controllers = array('twofactor', 'user');
-        $actions = array('code', 'check', 'logout');
+        $ignore = ($controller === 'twofactor' && in_array($action, array('code', 'check'))) || ($controller === 'user' && $action === 'logout');
 
-        if ($this->userSession->has2FA() && ! $this->userSession->check2FA() && ! in_array($controller, $controllers) && ! in_array($action, $actions)) {
+        if ($ignore === false && $this->userSession->has2FA() && ! $this->userSession->check2FA()) {
 
             if ($this->request->isAjax()) {
                 $this->response->text('Not Authorized', 401);
diff --git a/app/Controller/Twofactor.php b/app/Controller/Twofactor.php
index 7711666b..48954dc8 100644
--- a/app/Controller/Twofactor.php
+++ b/app/Controller/Twofactor.php
@@ -72,6 +72,9 @@ class Twofactor extends User
             ));
         }
 
+        // Allow the user to test or disable the feature
+        $this->session['user']['twofactor_activated'] = false;
+
         $this->session->flash(t('User updated successfully.'));
         $this->response->redirect($this->helper->url('twofactor', 'index', array('user_id' => $user['id'])));
     }