From 60f40f10eeeb95443dea9d964660dc477c0bd8ea Mon Sep 17 00:00:00 2001
From: Matt Westcott <matt@west.co.tt>
Date: Wed, 7 Jun 2017 23:48:06 +0100
Subject: [PATCH] Move 'mark view restriction as passed' logic to the
 BaseViewRestriction model

---
 wagtail/wagtailcore/models.py | 15 +++++++++++++++
 wagtail/wagtailcore/views.py  | 11 +----------
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/wagtail/wagtailcore/models.py b/wagtail/wagtailcore/models.py
index aeb024f041..1763bad2bd 100644
--- a/wagtail/wagtailcore/models.py
+++ b/wagtail/wagtailcore/models.py
@@ -1908,6 +1908,21 @@ class BaseViewRestriction(models.Model):
 
         return True
 
+    def mark_as_passed(self, request):
+        """
+        Update the session data in the request to mark the user as having passed this
+        view restriction
+        """
+        has_existing_session = (settings.SESSION_COOKIE_NAME in request.COOKIES)
+        passed_restrictions = request.session.setdefault(self.passed_view_restrictions_session_key, [])
+        if self.id not in passed_restrictions:
+            passed_restrictions.append(self.id)
+            request.session[self.passed_view_restrictions_session_key] = passed_restrictions
+        if not has_existing_session:
+            # if this is a session we've created, set it to expire at the end
+            # of the browser session
+            request.session.set_expiry(0)
+
     class Meta:
         abstract = True
         verbose_name = _('view restriction')
diff --git a/wagtail/wagtailcore/views.py b/wagtail/wagtailcore/views.py
index 4cfc764eda..1346d4cdd8 100644
--- a/wagtail/wagtailcore/views.py
+++ b/wagtail/wagtailcore/views.py
@@ -1,6 +1,5 @@
 from __future__ import absolute_import, unicode_literals
 
-from django.conf import settings
 from django.core.urlresolvers import reverse
 from django.http import Http404, HttpResponse
 from django.shortcuts import get_object_or_404, redirect
@@ -38,15 +37,7 @@ def authenticate_with_password(request, page_view_restriction_id, page_id):
     if request.method == 'POST':
         form = PasswordPageViewRestrictionForm(request.POST, instance=restriction)
         if form.is_valid():
-            has_existing_session = (settings.SESSION_COOKIE_NAME in request.COOKIES)
-            passed_restrictions = request.session.setdefault('passed_page_view_restrictions', [])
-            if restriction.id not in passed_restrictions:
-                passed_restrictions.append(restriction.id)
-                request.session['passed_page_view_restrictions'] = passed_restrictions
-            if not has_existing_session:
-                # if this is a session we've created, set it to expire at the end
-                # of the browser session
-                request.session.set_expiry(0)
+            restriction.mark_as_passed(request)
 
             return redirect(form.cleaned_data['return_url'])
     else: