Add collection view restriction model

wagtail/wagtailcore/migrations/0037_collectionviewrestriction.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.10.4 on 2016-12-19 15:32
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('wagtailcore', '0036_populate_page_last_published_at'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='CollectionViewRestriction',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('restriction_type', models.CharField(choices=[('none', 'Public'), ('login', 'Private, accessible to logged-in users'), ('password', 'Private, accessible with the following password'), ('groups', 'Private, accessible to users in specific groups')], max_length=20)),
+                ('password', models.CharField(blank=True, max_length=255, verbose_name='password')),
+                ('collection', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='view_restrictions', to='wagtailcore.Collection', verbose_name='collection')),
+                ('groups', models.ManyToManyField(blank=True, to='auth.Group', verbose_name='groups')),
+            ],
+            options={
+                'verbose_name': 'collection view restriction',
+                'verbose_name_plural': 'collection view restrictions',
+            },
+        ),
+        migrations.AlterField(
+            model_name='pageviewrestriction',
+            name='groups',
+            field=models.ManyToManyField(blank=True, to='auth.Group', verbose_name='groups'),
+        ),
+    ]

wagtail/wagtailcore/models.py

@@ -1871,7 +1871,7 @@ class PagePermissionTester(object):
         return True
 
 
-class PageViewRestriction(models.Model):
+class BaseViewRestriction(models.Model):
     NONE = 'none'
     PASSWORD = 'password'
     GROUPS = 'groups'
@@ -1886,23 +1886,20 @@ class PageViewRestriction(models.Model):
 
     restriction_type = models.CharField(
         max_length=20, choices=RESTRICTION_CHOICES)
-    page = models.ForeignKey(
-        'Page', verbose_name=_('page'), related_name='view_restrictions', on_delete=models.CASCADE
-    )
     password = models.CharField(verbose_name=_('password'), max_length=255, blank=True)
-    groups = models.ManyToManyField(Group, blank=True)
+    groups = models.ManyToManyField(Group, verbose_name=_('groups'), blank=True)
 
     def accept_request(self, request):
-        if self.restriction_type == PageViewRestriction.PASSWORD:
-            passed_restrictions = request.session.get('passed_page_view_restrictions', [])
+        if self.restriction_type == BaseViewRestriction.PASSWORD:
+            passed_restrictions = request.session.get(self.passed_view_restrictions_session_key, [])
             if self.id not in passed_restrictions:
                 return False
 
-        elif self.restriction_type == PageViewRestriction.LOGIN:
+        elif self.restriction_type == BaseViewRestriction.LOGIN:
             if not user_is_authenticated(request.user):
                 return False
 
-        elif self.restriction_type == PageViewRestriction.GROUPS:
+        elif self.restriction_type == BaseViewRestriction.GROUPS:
             if not request.user.is_superuser:
                 current_user_groups = request.user.groups.all()
 
@@ -1911,6 +1908,19 @@ class PageViewRestriction(models.Model):
 
         return True
 
+    class Meta:
+        abstract = True
+        verbose_name = _('view restriction')
+        verbose_name_plural = _('view restrictions')
+
+
+class PageViewRestriction(BaseViewRestriction):
+    page = models.ForeignKey(
+        'Page', verbose_name=_('page'), related_name='view_restrictions', on_delete=models.CASCADE
+    )
+
+    passed_view_restrictions_session_key = 'passed_page_view_restrictions'
+
     class Meta:
         verbose_name = _('page view restriction')
         verbose_name_plural = _('page view restrictions')
@@ -1924,6 +1934,21 @@ class BaseCollectionManager(models.Manager):
 CollectionManager = BaseCollectionManager.from_queryset(TreeQuerySet)
 
 
+class CollectionViewRestriction(BaseViewRestriction):
+    collection = models.ForeignKey(
+        'Collection',
+        verbose_name=_('collection'),
+        related_name='view_restrictions',
+        on_delete=models.CASCADE
+    )
+
+    passed_view_restrictions_session_key = 'passed_collection_view_restrictions'
+
+    class Meta:
+        verbose_name = _('collection view restriction')
+        verbose_name_plural = _('collection view restrictions')
+
+
 @python_2_unicode_compatible
 class Collection(MP_Node):
     """
@@ -1951,6 +1976,10 @@ class Collection(MP_Node):
     def get_prev_siblings(self, inclusive=False):
         return self.get_siblings(inclusive).filter(path__lte=self.path).order_by('-path')
 
+    def get_view_restrictions(self):
+        """Return a query set of all collection view restrictions that apply to this collection"""
+        return CollectionViewRestriction.objects.filter(collection__in=self.get_ancestors(inclusive=True))
+
     class Meta:
         verbose_name = _('collection')
         verbose_name_plural = _('collections')