Merge lp:~canonical-isd-hackers/canonical-identity-provider/enforce-user-requirestwo-factor into lp:canonical-identity-provider/release

[Michael Foord]

Adds new sso_login_required decorator to wrap up django's login_required and add twofactor checks

=== modified file 'identityprovider/decorators.py'
--- identityprovider/decorators.py	2012-01-24 12:23:06 +0000
+++ identityprovider/decorators.py	2012-02-09 22:30:07 +0000
@@ -2,9 +2,17 @@
 # GNU Affero General Public License version 3 (see the file LICENSE).
 
 import logging
+from functools import wraps
 
 from django.template import RequestContext
 from django.http import Http404, HttpResponseRedirect, HttpResponseForbidden
+from django.contrib.auth import REDIRECT_FIELD_NAME
+from django.utils.decorators import available_attrs
+from django.contrib.auth.decorators import ( 
+    login_required as django_login_required,
+)
+from django.contrib.auth.views import redirect_to_login
+from django.core.urlresolvers import reverse
 from django.core.cache import cache
 from django.conf import settings
 from django.template.loader import render_to_string
@@ -16,6 +24,8 @@
 
 from identityprovider import auth
 from identityprovider.cookies import set_test_cookie, test_cookie_worked
+from identityprovider.models import twofactor
+
 
 def guest_required(func):
     @functools.wraps(func)
@@ -38,6 +48,33 @@
     return _dont_cache_decorator
 
 
+def sso_login_required(function=None, redirect_field_name=REDIRECT_FIELD_NAME, login_url=None):
+    """
+    Decorator the wraps up django's login_required, and also checks for 2f
+    """
+    def decorator(view_func):
+        @wraps(view_func, assigned=available_attrs(view_func))
+        def _wrapped_view(request, *args, **kwargs):
+            if (twofactor.user_requires_twofactor_auth(request, request.user) and
+                    not twofactor.is_authenticated(request)):
+                # some views send token as None, which breaks reverse
+                if 'token' in kwargs and kwargs['token'] == None:
+                    del kwargs['token']
+                return redirect_to_login(
+                        request.get_full_path(),
+                        reverse('twofactor', kwargs=kwargs),
+                        redirect_field_name)
+            return view_func(request, *args, **kwargs)
+        return django_login_required(
+                _wrapped_view, 
+                redirect_field_name,
+                login_url)
+
+    if function:
+        return decorator(function)
+    return decorator
+
+
 def check_readonly(func):
     """ A readonly aware decorator.
 

=== modified file 'identityprovider/models/twofactor.py'
--- identityprovider/models/twofactor.py	2012-01-30 14:42:31 +0000
+++ identityprovider/models/twofactor.py	2012-02-09 22:30:07 +0000
@@ -1,10 +1,15 @@
 # Copyright 2010-2012 Canonical Ltd.  This software is licensed under the
 # GNU Affero General Public License version 3 (see the file LICENSE).
+
+import logging
+
 from oath.hotp import accept_hotp
 from django.db import models
 from django.conf import settings
 from django.utils.translation import ugettext as _
 
+from gargoyle import gargoyle
+
 from identityprovider.models import Account
 from identityprovider.utils import AuthenticationError
 
@@ -20,6 +25,37 @@
     else:
         return None
 
+TWOFACTOR_KEY = 'logged_in_via_two_factor'
+
+def is_twofactor_enabled(request):
+    return gargoyle.is_active('TWOFACTOR', request)
+
+def login(request):
+    request.session[TWOFACTOR_KEY] = True
+
+def logout(request):
+    if TWOFACTOR_KEY in request.session:
+        del request.session[TWOFACTOR_KEY]
+
+def is_authenticated(request):
+    return request.session.get(TWOFACTOR_KEY, False)
+
+def site_requires_twofactor_auth(request, token):
+    if not is_twofactor_enabled(request):
+        return False
+    # TODO check if site requires 2f (pape?)
+    return False
+
+def user_requires_twofactor_auth(request, account):
+    if not is_twofactor_enabled(request):
+        return False
+    if account.twofactor_required:
+        if account.has_twofactor_devices():
+            return True
+        logging.warning('Account %s requires two-factor but has no devices' % account.openid_identifier)
+    return False
+
+
 class AuthenticationDevice(models.Model):
     account = models.ForeignKey(Account, related_name='devices')
     key = models.TextField()

=== added file 'identityprovider/tests/acceptance/devices/__init__.py'
=== modified file 'identityprovider/tests/acceptance/devices/add_device.py'
--- identityprovider/tests/acceptance/devices/add_device.py	2012-01-27 16:28:56 +0000
+++ identityprovider/tests/acceptance/devices/add_device.py	2012-02-09 22:30:07 +0000
@@ -4,7 +4,7 @@
 from sst.actions import *
 
 from identityprovider.tests.acceptance import _helpers as helpers
-from _helpers import click_add_device_button, click_add_new_device_link
+from identityprovider.tests.acceptance.devices._helpers import click_add_device_button, click_add_new_device_link
 
 
 config.set_base_url_from_env()

=== modified file 'identityprovider/tests/acceptance/devices/login.py'
--- identityprovider/tests/acceptance/devices/login.py	2012-02-03 13:57:54 +0000
+++ identityprovider/tests/acceptance/devices/login.py	2012-02-09 22:30:07 +0000
@@ -4,7 +4,7 @@
 from canonical.isd.tests.sst import config
 from sst.actions import *
 
-from _helpers import add_device
+from identityprovider.tests.acceptance.devices._helpers import add_device
 from identityprovider.tests.acceptance import _helpers as helpers
 
 def enter_otp(otp):

=== modified file 'identityprovider/tests/acceptance/devices/prefs.py'
--- identityprovider/tests/acceptance/devices/prefs.py	2012-01-27 16:28:56 +0000
+++ identityprovider/tests/acceptance/devices/prefs.py	2012-02-09 22:30:07 +0000
@@ -1,8 +1,9 @@
 # Test user two factor preferences
+from oath import hotp
 from canonical.isd.tests.sst import config
 from sst.actions import *
 
-from _helpers import add_device
+from identityprovider.tests.acceptance.devices._helpers import add_device
 from identityprovider.tests.acceptance import _helpers as helpers
 
 
@@ -26,7 +27,7 @@
                disabled="disabled")
 
 # Add an authentication device
-add_device()
+aes_key = add_device()
 
 # Test preferences form is now enabled
 assert_attribute(get_element(tag="input", name="two-factor", value="always"),
@@ -45,7 +46,16 @@
 set_radio_value(get_element(tag="input", name="two-factor", value="always"))
 click_button("authentication-prefs")
 
+# so we now should now have the twofactor auth form
+assert_url('/two_factor_auth?next=/device-list')
+
+# 2 factor login
+otp = hotp.hotp(aes_key, 1)
+write_textfield('id_oath_token', otp)
+click_button(get_element(type='submit'))
+
 # Check that "always" is now shown in the user preferences
+assert_url('/device-list')
 assert_attribute(get_element(tag="input", name="two-factor", value="always"),
                  'checked', 'true')
 assert_attribute(get_element(tag="input", name="two-factor", value="as-needed"),

=== modified file 'identityprovider/tests/acceptance/devices/remove_device.py'
--- identityprovider/tests/acceptance/devices/remove_device.py	2012-01-27 16:29:24 +0000
+++ identityprovider/tests/acceptance/devices/remove_device.py	2012-02-09 22:30:07 +0000
@@ -3,7 +3,7 @@
 from sst.actions import *
 
 from identityprovider.tests.acceptance import _helpers as helpers
-from _helpers import add_device
+from identityprovider.tests.acceptance.devices._helpers import add_device
 
 
 def click_remove_button():

=== modified file 'identityprovider/tests/acceptance/devices/rename_device.py'
--- identityprovider/tests/acceptance/devices/rename_device.py	2012-01-27 16:28:56 +0000
+++ identityprovider/tests/acceptance/devices/rename_device.py	2012-02-09 22:30:07 +0000
@@ -2,7 +2,7 @@
 from sst.actions import *
 
 from identityprovider.tests.acceptance import _helpers as helpers
-from _helpers import (
+from identityprovider.tests.acceptance.devices._helpers import (
     add_device,
     assert_device,
     assert_no_device,

=== added file 'identityprovider/tests/acceptance/views_protected.csv'
--- identityprovider/tests/acceptance/views_protected.csv	1970-01-01 00:00:00 +0000
+++ identityprovider/tests/acceptance/views_protected.csv	2012-02-09 22:30:07 +0000
@@ -0,0 +1,8 @@
+url
+'/+edit'
+'/+emails'
+'/+index'
+'/+applications'
+'/+new-email'
+'/+verify-email'
+'/+remove-email'

=== added file 'identityprovider/tests/acceptance/views_protected.py'
--- identityprovider/tests/acceptance/views_protected.py	1970-01-01 00:00:00 +0000
+++ identityprovider/tests/acceptance/views_protected.py	2012-02-09 22:30:07 +0000
@@ -0,0 +1,44 @@
+import sys
+from urllib import quote
+
+from canonical.isd.tests.sst import config
+from sst.actions import *
+
+import _helpers as helpers
+from devices import _helpers as dev_helpers
+
+config.set_base_url_from_env()
+
+try:
+    import _views_protected_shared as shared
+except ImportError:
+    class module(object):
+        pass
+    m = module()
+
+    m.twofactor_password = 'Admin009'
+    m.twofactor_email = helpers.register_account(password=m.twofactor_password)
+    m.aes_key = dev_helpers.add_device()
+    # Change the preference to always require 2 factor authentication and save
+    set_radio_value(get_element(tag="input", name="two-factor", value="always"))
+    click_button("authentication-prefs")
+    assert_url('/two_factor_auth?next=/device-list')
+    helpers.logout()
+
+    sys.modules['_views_protected_shared'] = m
+    import _views_protected_shared as shared
+
+# are we redirected?
+go_to(url)
+wait_for(assert_title, 'Log in')
+
+helpers.login(shared.twofactor_email, shared.twofactor_password)
+
+go_to(url)
+# should be redirected
+assert_url('/two_factor_auth?next=%s' % quote(url))
+
+helpers.logout()
+
+
+

=== modified file 'identityprovider/tests/unit/test_decorators.py'
--- identityprovider/tests/unit/test_decorators.py	2012-01-10 21:17:33 +0000
+++ identityprovider/tests/unit/test_decorators.py	2012-02-09 22:30:07 +0000
@@ -1,9 +1,9 @@
 from datetime import datetime
 from django.test import TestCase
 
-from mock import Mock, patch
+from mock import Mock, patch, MagicMock
 
-from identityprovider.decorators import ratelimit
+from identityprovider.decorators import ratelimit, sso_login_required
 
 
 class RatelimitTestCase(TestCase):
@@ -32,3 +32,75 @@
             limiter.current_key(request)
 
         self.assertEqual(mock_datetime.utcnow.called, True)
+
+
+class SSOLoginRequiredTestCase(TestCase):
+
+    @staticmethod
+    def fake_view(request, *args, **kwargs):
+        return "SUCCESS"
+
+    def mock_request(self, authed, path='/'):
+        mock_request = MagicMock()
+        mock_request.method = 'GET'
+        mock_request.user.is_authenticated.return_value = authed
+        mock_request.POST = {}
+        mock_request.META = {}
+        mock_request.build_absolute_uri.return_value = path
+        mock_request.get_full_path.return_value = path
+        return mock_request
+
+    def test_django_login_check_still_works(self):
+        view = sso_login_required(self.fake_view, login_url='/login')
+        response = view(self.mock_request(False, '/target'))
+        self.assertEqual(response.status_code, 302)
+        self.assertEqual(response['Location'], '/login?next=/target')
+
+    @patch('identityprovider.decorators.twofactor.user_requires_twofactor_auth')
+    def test_allowed_2f_not_required(self, mock_user):
+        mock_user.return_value = False
+        view = sso_login_required(self.fake_view, login_url='/login')
+        response = view(self.mock_request(True))
+        self.assertEqual(response, 'SUCCESS')
+
+    @patch('identityprovider.decorators.twofactor.is_authenticated')
+    @patch('identityprovider.decorators.twofactor.user_requires_twofactor_auth')
+    def test_allowed_when_2f_required_and_2f_authed(self, mock_user, mock_auth):
+        mock_user.return_value = True
+        mock_auth.return_value = True
+        view = sso_login_required(self.fake_view, login_url='/login')
+        response = view(self.mock_request(True))
+        self.assertEqual(response, 'SUCCESS')
+
+    @patch('identityprovider.decorators.twofactor.is_authenticated')
+    @patch('identityprovider.decorators.twofactor.user_requires_twofactor_auth')
+    def test_redirected_when_2f_required_and_not_2f_authed(self, mock_user, mock_auth):
+        mock_user.return_value = True
+        mock_auth.return_value = False
+        view = sso_login_required(self.fake_view, login_url='/login')
+        response = view(self.mock_request(True, '/target'))
+        self.assertEqual(response.status_code, 302)
+        self.assertEqual(response['Location'], '/two_factor_auth?next=/target')
+
+    def test_decoration_no_params(self):
+        @sso_login_required
+        def view(request, *args, **kwargs):
+            return "SUCCESS"
+        response = view(self.mock_request(True, '/target'))
+        # if this works, the decorator handles this invocation properly
+        self.assertEqual(response, 'SUCCESS')
+
+    def test_decoration_with_paramss(self):
+        @sso_login_required(login_url='/alt_login')
+        def view(request, *args, **kwargs):
+            return "SUCCESS"
+        response = view(self.mock_request(True, '/target'))
+        # if this works, the decorator handles this invocation properly
+        self.assertEqual(response, 'SUCCESS')
+
+
+
+
+
+
+

=== modified file 'identityprovider/tests/unit/test_views_ui.py'
--- identityprovider/tests/unit/test_views_ui.py	2012-02-09 11:47:25 +0000
+++ identityprovider/tests/unit/test_views_ui.py	2012-02-09 22:30:07 +0000
@@ -30,7 +30,7 @@
 from identityprovider.models import EmailAddress, Person, OpenIDRPConfig
 from identityprovider.models.const import EmailStatus
 from identityprovider.views import ui, server
-from identityprovider.models import Account, AuthToken
+from identityprovider.models import Account, AuthToken, twofactor
 from identityprovider.models.authtoken import (
     InvalidEmailException,
     create_token,
@@ -1180,14 +1180,14 @@
             data['next'] = next
         return self.client.post('/two_factor_auth', data)
 
-    @patch('identityprovider.views.ui.site_requires_twofactor_auth')
+    @patch('identityprovider.views.ui.twofactor.site_requires_twofactor_auth')
     def test_when_site_does_not_require_twofactor_no_oath_field(self, mock_site):
         mock_site.return_value = False
         response = self.client.get('/+login')
         tree = PyQuery(fromstring(response.content))
         self.assertEqual(len(tree.find('input[name="oath_token"]')), 0)
 
-    @patch('identityprovider.views.ui.site_requires_twofactor_auth')
+    @patch('identityprovider.views.ui.twofactor.site_requires_twofactor_auth')
     def test_when_site_requires_twofactor_oath_field_is_present(self, mock_site):
         mock_site.return_value = True
         response = self.client.get('/+login')
@@ -1197,7 +1197,7 @@
         oath = inputs[0]
         self.assertEqual(oath.attrib['autocomplete'], 'off')
 
-    @patch('identityprovider.views.ui.site_requires_twofactor_auth')
+    @patch('identityprovider.views.ui.twofactor.site_requires_twofactor_auth')
     def test_when_site_requires_twofactor_oath_field_is_present_lp(self, mock_site):
         mock_site.return_value = True
         response = self.client.get('/+login')
@@ -1207,27 +1207,26 @@
         oath = inputs[0]
         self.assertEqual(oath.attrib['autocomplete'], 'off')
 
-
-    @patch('identityprovider.views.ui.user_requires_twofactor_auth')
+    @patch('identityprovider.views.ui.twofactor.user_requires_twofactor_auth')
     def test_when_user_requires_twofactor_redirected(self, mock_user):
         mock_user.return_value = True
         r = self.do_login()
         self.assertRedirects(r, '/two_factor_auth')
 
-    @patch('identityprovider.views.ui.user_requires_twofactor_auth')
+    @patch('identityprovider.views.ui.twofactor.user_requires_twofactor_auth')
     def test_when_user_requires_twofactor_redirected_with_token(self, mock_user):
         mock_user.return_value = True
         token = 'a' * 16
         r = self.do_login(token=token)
         self.assertRedirects(r, '/%s/two_factor_auth' % token)
 
-    @patch('identityprovider.views.ui.user_requires_twofactor_auth')
+    @patch('identityprovider.views.ui.twofactor.user_requires_twofactor_auth')
     def test_when_user_requires_twofactor_redirected_with_next(self, mock_user):
         mock_user.return_value = True
         r = self.do_login(next='/+edit')
         self.assertRedirects(r, '/two_factor_auth?next=/%2Bedit')
 
-    @patch('identityprovider.views.ui.site_requires_twofactor_auth')
+    @patch('identityprovider.views.ui.twofactor.site_requires_twofactor_auth')
     @patch('identityprovider.views.ui.authenticate_device')
     def test_full_twofactor_login_redirected_to_next(self, mock_auth, mock_site):
         mock_auth.return_value = True
@@ -1236,7 +1235,7 @@
         mock_auth.assert_called_once_with(self.account, '123456')
         self.assertRedirects(r, '/+edit')
 
-    @patch('identityprovider.views.ui.site_requires_twofactor_auth')
+    @patch('identityprovider.views.ui.twofactor.site_requires_twofactor_auth')
     def test_full_twofactor_login_illformed_token_displays_error(self, mock_site):
         mock_site.return_value = True
         response = self.do_login(next='/+edit', oath_token='XXX')
@@ -1246,7 +1245,7 @@
         self.assertEqual(len(error), 1)
         self.assertEqual(error[0].text.strip(), "Please enter a 6-digit or 8-digit one-time password.")
 
-    @patch('identityprovider.views.ui.site_requires_twofactor_auth')
+    @patch('identityprovider.views.ui.twofactor.site_requires_twofactor_auth')
     @patch('identityprovider.views.ui.authenticate_device')
     def test_full_twofactor_login_invalid_code_displays_error(self, mock_auth, mock_site):
         mock_auth.side_effect = AuthenticationError("ERRORMSG")
@@ -1314,12 +1313,12 @@
     def test_user_requires_twofactor_auth(self):
         def f(require2f, has_devices, returns, logs):
             with patch('logging.warning') as mock_warning:
-                with patch('identityprovider.views.ui.is_twofactor_enabled') as mock_flag:
+                with patch('identityprovider.views.ui.twofactor.is_twofactor_enabled') as mock_flag:
                     mock_flag.return_value = True
                     account = Mock()
                     account.twofactor_required = require2f
                     account.has_twofactor_devices.return_value = has_devices
-                    r = ui.user_requires_twofactor_auth(None, account)
+                    r = twofactor.user_requires_twofactor_auth(None, account)
                     self.assertEqual(returns, r)
                     self.assertEqual(logs, mock_warning.called)
         f(require2f=False, has_devices=False, returns=False, logs=False)

=== modified file 'identityprovider/views/account.py'
--- identityprovider/views/account.py	2012-01-13 14:49:47 +0000
+++ identityprovider/views/account.py	2012-02-09 22:30:07 +0000
@@ -4,7 +4,6 @@
 from openid.yadis.constants import YADIS_HEADER_NAME, YADIS_CONTENT_TYPE
 
 from django.conf import settings
-from django.contrib.auth.decorators import login_required
 from django.http import HttpResponseRedirect
 from django.shortcuts import render_to_response, get_object_or_404
 from django.template import RequestContext
@@ -12,7 +11,7 @@
 from django.views.decorators.vary import vary_on_headers
 
 from identityprovider.cookies import set_test_cookie, test_cookie_worked
-from identityprovider.decorators import check_readonly
+from identityprovider.decorators import check_readonly, sso_login_required
 from identityprovider.forms import EditAccountForm, LoginForm, NewEmailForm
 from identityprovider.models import (send_validation_email_request,
     EmailAddress)
@@ -58,7 +57,7 @@
     return set_test_cookie(render_to_response('cookies.html'))
 
 
-@login_required
+@sso_login_required
 def account_index(request, token=None):
     request.token = token
     account = request.user
@@ -86,7 +85,7 @@
     return render_to_response('account/edit.html', context)
 
 
-@login_required
+@sso_login_required
 def account_emails(request):
     account = request.user
     form = NewEmailForm()
@@ -103,7 +102,7 @@
     return render_to_response('account/emails.html', context)
 
 
-@login_required
+@sso_login_required
 def account_deactivate(request, token=None):
     request.token = token
     import django.contrib.auth as auth
@@ -152,7 +151,7 @@
     )
 
 
-@login_required
+@sso_login_required
 @check_readonly
 def new_email(request, token=None, emailid=None):
     request.token = token
@@ -171,7 +170,7 @@
     return render_to_response('account/new_email.html', context)
 
 
-@login_required
+@sso_login_required
 def verify_email(request, token=None):
     request.token = token
     emailid = request.GET.get('id', 0)
@@ -180,7 +179,7 @@
     return _send_verification_email(request, email.email, token)
 
 
-@login_required
+@sso_login_required
 def delete_email(request, token=None):
     request.token = token
     emailid = request.GET.get('id', 0)
@@ -201,7 +200,7 @@
         return render_to_response('account/delete_email.html', context)
 
 
-@login_required
+@sso_login_required
 @check_readonly
 def applications(request):
     if request.method == 'POST':

=== modified file 'identityprovider/views/devices.py'
--- identityprovider/views/devices.py	2012-02-01 22:34:09 +0000
+++ identityprovider/views/devices.py	2012-02-09 22:30:07 +0000
@@ -4,7 +4,6 @@
 import re
 from base64 import b16encode
 
-from django.contrib.auth.decorators import login_required
 from django.core.urlresolvers import reverse
 from django.conf import settings
 from django.shortcuts import get_object_or_404, render_to_response
@@ -18,6 +17,7 @@
 from identityprovider.models import AuthenticationDevice
 from identityprovider.models.twofactor import get_otp_type
 from identityprovider.forms import HOTPDeviceForm, DeviceRenameForm
+from identityprovider.decorators import sso_login_required
 
 from .device_urldata import *
 from .utils import HttpResponseRedirect, HttpResponseSeeOther, allow_only
@@ -52,7 +52,7 @@
     return rand_bytes(n)
 
 
-@login_required
+@sso_login_required
 @require_twofactor
 @allow_only('GET', 'POST')
 def device_list(request):
@@ -88,7 +88,7 @@
     return render_to_response('device/list.html', context)
 
 
-@login_required
+@sso_login_required
 @require_twofactor
 @allow_only('GET', 'POST')
 def device_addition(request):
@@ -173,7 +173,7 @@
 
 
 
-@login_required
+@sso_login_required
 @require_twofactor
 @allow_only('GET', 'POST')
 def device_verification(request, device_id):
@@ -187,7 +187,7 @@
     pass
 
 
-@login_required
+@sso_login_required
 @require_twofactor
 @allow_only('GET', 'POST')
 def device_removal(request, device_id):
@@ -235,10 +235,10 @@
                     device_list_path=reverse(DEVICE_LIST),
                     form=form))
 
-device_rename = login_required(DeviceRenameView.as_view())
-
-
-@login_required
+device_rename = sso_login_required(DeviceRenameView.as_view())
+
+
+@sso_login_required
 @require_twofactor
 @allow_only('GET', 'POST')
 def pad_removal(request):
@@ -248,7 +248,7 @@
     return HttpResponseSeeOther('/device-list')
 
 
-@login_required
+@sso_login_required
 @require_twofactor
 @allow_only('GET', 'POST')
 def device_lost(request):

=== modified file 'identityprovider/views/server.py'
--- identityprovider/views/server.py	2012-01-24 12:23:06 +0000
+++ identityprovider/views/server.py	2012-02-09 22:30:07 +0000
@@ -33,7 +33,6 @@
 import django.contrib.auth as auth
 
 from django.conf import settings
-from django.contrib.auth.decorators import login_required
 from django.http import (
     Http404,
     HttpResponse,
@@ -55,6 +54,7 @@
 import identityprovider.signed as signed
 
 from identityprovider.const import LAUNCHPAD_TEAMS_NS
+from identityprovider.decorators import sso_login_required
 from identityprovider.forms import (
     PreAuthorizeForm,
     SRegRequestForm,
@@ -532,7 +532,7 @@
         openid_response.addExtension(sreg_response)
 
 
-@login_required
+@sso_login_required
 def _process_decide(request, orequest, decision):
     oresponse = orequest.answer(decision,
         identity=request.user.openid_identity_url)

=== modified file 'identityprovider/views/ui.py'
--- identityprovider/views/ui.py	2012-02-06 16:54:13 +0000
+++ identityprovider/views/ui.py	2012-02-09 22:30:07 +0000
@@ -24,12 +24,16 @@
 from django.utils.decorators import method_decorator
 from django import forms
 
-from gargoyle import gargoyle
 from gargoyle.decorators import switch_is_active
 
 from identityprovider.branding import current_brand
-from identityprovider.decorators import (check_readonly, dont_cache,
-    guest_required, limitlogin, requires_cookies)
+from identityprovider.decorators import (
+        check_readonly,
+        dont_cache,
+        guest_required,
+        limitlogin,
+        requires_cookies,
+    )
 from identityprovider.forms import (
     ConfirmNewAccountForm, 
     ForgotPasswordForm, 
@@ -48,6 +52,7 @@
 from identityprovider.models.const import (AccountStatus, EmailStatus,
     LoginTokenType)
 from identityprovider.models.openidmodels import OpenIDRPConfig
+from identityprovider.models import twofactor
 import identityprovider.signed as signed
 from identityprovider.utils import (
     AuthenticationError,
@@ -62,7 +67,6 @@
 from identityprovider.views.utils import get_rpconfig, set_session_token_info
 
 
-
 logger = logging.getLogger('sso')
 
 
@@ -91,30 +95,6 @@
         errors.append(str(exc))
 
 
-def is_twofactor_enabled(request):
-    return gargoyle.is_active('TWOFACTOR', request)
-
-
-# these may need to live elsewhere, but putting here for now
-def site_requires_twofactor_auth(request, token):
-    if not is_twofactor_enabled(request):
-        return False
-    # TODO check if site requires 2f (pape?)
-    return False
-
-
-def user_requires_twofactor_auth(request, account):
-    if not is_twofactor_enabled(request):
-        return False
-    if account.twofactor_required:
-        if account.has_twofactor_devices():
-            return True
-        logging.warning('Account %s requires two-factor but has no devices' % account.openid_identifier)
-    return False
-
-
-
-
 def authenticate_user(email, password):
     """Attempts to authenticate a user. Returns account on success, or throws
     AuthenticationErrors on failure"""
@@ -191,12 +171,12 @@
     template_name = 'registration/login.html'
 
     def get_login_type(self, request, token):
-        twofactor = site_requires_twofactor_auth(request, token)
-        return twofactor, TwoFactorLoginForm if twofactor else LoginForm
+        required = twofactor.site_requires_twofactor_auth(request, token)
+        return required, TwoFactorLoginForm if required else LoginForm
 
     def get(self, request, token=None, rpconfig=None):
         rpconfig = self.setup(request, token, rpconfig)
-        twofactor, form_cls = self.get_login_type(request, token)
+        required2f, form_cls = self.get_login_type(request, token)
         # track login attempts
         stats.increment('flows.login', key='requested', rpconfig=rpconfig)
         return self.render(request, token, rpconfig, form_cls())
@@ -223,7 +203,8 @@
         # handle case where the user's account requires two factor but we
         # didn't know that until we authenticated them
         oath_token = form.cleaned_data.get('oath_token', None)
-        if not oath_token and user_requires_twofactor_auth(request, account):
+        if (not oath_token and
+                twofactor.user_requires_twofactor_auth(request, account)):
             kwargs = {'token':token} if token else {}
             url = reverse('twofactor', kwargs=kwargs)
             if next and _is_safe_redirect_url(next):
@@ -234,7 +215,7 @@
         if site_twofactor or oath_token:
             try:
                 authenticate_device(account, oath_token)
-                request.session['logged_in_via_two_factor'] = True
+                twofactor.login(request)
             except AuthenticationError as e:
                 return self.display_errors(request, token, rpconfig, form, e)
 
@@ -279,7 +260,7 @@
 
         try:
             authenticate_device(account, form.cleaned_data['oath_token'])
-            request.session['logged_in_via_two_step'] = True
+            twofactor.login(request)
             stats.increment('flows.login', key='success', rpconfig=rpconfig)
         except AuthenticationError as e:
             return self.display_errors(request, token, rpconfig, form, e)
@@ -297,6 +278,7 @@
 
     @method_decorator(dont_cache)
     @method_decorator(limitlogin())
+    # for this page, we need to be logged in but NOT twofactored
     @method_decorator(login_required)
     @method_decorator(switch_is_active('TWOFACTOR'))
     def dispatch(self, *args, **kwargs):
@@ -380,6 +362,7 @@
         # We don't want to lose session[token] when we log the user out
         raw_orequest = request.session.get(token, None)
         auth.logout(request)
+        twofactor.logout(request)
         self.set_orequest(request.session, token, raw_orequest)
 
         template_file = ('%s/registration/logout.html' %