add role permission

app/decorators.py

@@ -0,0 +1,17 @@
+from functools import wraps
+from flask import abort
+from flask_login import current_user
+from .models import Permission
+
+def permission_required(permission):
+    def decorator(f):
+        @wraps(f)
+        def decorated_function(*args, **kwargs):
+            if not current_user.can(permission):
+                abort(403)
+            return f(*args, **kwargs)
+        return decorated_function
+    return decorator
+
+def admin_required(f):
+    return permission_required(Permission.ADMIN)(f)

app/main/__init__.py

@@ -1,3 +1,9 @@
 from flask import Blueprint
 main = Blueprint('main', __name__)
-from . import views, forms
+from . import views, forms
+from ..models import Permission
+
+
+@main.app_context_processor
+def inject_permissions():
+    return dict(Permission=Permission)

app/models.py

@@ -1,20 +1,74 @@
 from . import db
 from werkzeug.security import generate_password_hash, check_password_hash
-from flask_login import UserMixin
+from flask_login import UserMixin, AnonymousUserMixin
 from . import login_manager
 from flask import current_app
 from itsdangerous import TimedJSONWebSignatureSerializer as Serializer
 
 
+class Permission:
+    FOLLOW = 1
+    COMMENT = 2
+    WRITE = 4
+    MODERATE = 8
+    ADMIN = 16
+
+
 class Role(db.Model):
     __tablename__ = 'roles'
     id = db.Column(db.Integer, primary_key=True)
     name = db.Column(db.String(64), unique=True)
     users = db.relationship('User', backref="role", lazy='dynamic')
+    default = db.Column(db.Boolean, default=False, index=True)
+    permissions = db.Column(db.Integer)
+
+    def __init__(self, **kwargs):
+        super(Role, self).__init__(**kwargs)
+        if self.permissions is None:
+            self.permissions = 0
 
     def __repr__(self):
         return f"<Role {self.name}>"
 
+    def add_permission(self, perm):
+        if not self.has_permission(perm):
+            self.permissions += perm
+
+    def remove_permission(self, perm):
+        if self.has_permission(perm):
+            self.permissions -= perm
+
+    def reset_permissions(self):
+        self.permissions = 0
+
+    def has_permission(self, perm):
+        return self.permissions & perm == perm
+
+    @staticmethod
+    def insert_roles():
+        roles = {
+            'User': [Permission.FOLLOW, Permission.COMMENT, Permission.WRITE],
+            'Moderator': [
+                Permission.FOLLOW, Permission.COMMENT, Permission.WRITE,
+                Permission.MODERATE
+            ],
+            'Administrator': [
+                Permission.FOLLOW, Permission.COMMENT, Permission.WRITE,
+                Permission.MODERATE, Permission.ADMIN
+            ]
+        }
+        default_role = 'User'
+        for r in roles:
+            role = Role.query.filter_by(name=r).first()
+            if role is None:
+                role = Role(name=r)
+            role.reset_permissions()
+            for perm in roles[r]:
+                role.add_permission(perm)
+            role.default = (role.name == default_role)
+            db.session.add(role)
+        db.session.commit()
+
 
 class User(UserMixin, db.Model):
     __tablename__ = 'users'
@@ -25,6 +79,14 @@ class User(UserMixin, db.Model):
     email = db.Column(db.String(64), unique=True, index=True)
     confirmed = db.Column(db.Boolean, default=False)
 
+    def __init__(self, **kwargs):
+        super(User, self).__init__(**kwargs)
+        if self.role is None:
+            if self.email == current_app.config['FLASKY_ADMIN']:
+                self.role = Role.query.filter_by(name='Administrator').first()
+            if self.role is None:
+                self.role = Role.query.filter_by(default=True).first()
+
     @property
     def password(self):
         raise AttributeError("password is not a readable attribute")
@@ -55,6 +117,23 @@ def confirm(self, token):
     def __repr__(self):
         return f"<User {self.username}>"
 
+    def can(self, perm):
+        return self.role is not None and self.role.has_permission(perm)
+
+    def is_administrator(self):
+        return self.can(Permission.ADMIN)
+
+
+class AnonymousUser(AnonymousUserMixin):
+    def can(self, permissions):
+        return False
+
+    def is_administrator(self):
+        return False
+
+
+login_manager.anonymous_user = AnonymousUser
+
 
 @login_manager.user_loader
 def load_user(user_id):

config.py

@@ -13,7 +13,7 @@ class Config:
     MAIL_PASSWORD = os.environ.get('MAIL_PASSWORD')
     FLASKY_MAIL_SUBJECT_PREFIX = '[Flasky]'
     FLASKY_MAIL_SENDER = 'Flasky Admin <[email protected]>'
-    FLASKY_ADMIN = os.environ.get('FLASKY_ADMIN')
+    FLASKY_ADMIN = os.environ.get('FLASKY_ADMIN') or "[email protected]"
     SQLALCHEMY_TRACK_MODIFICATIONS = False
 
     @staticmethod

flasky.py

@@ -5,15 +5,15 @@
 from flask_migrate import Migrate
 
 from app import create_app, db
-from app.models import Role, User
+from app.models import Role, User, Permission
 
 app = create_app(os.environ.get('FLASK_CONFIG', 'default'))
 migrate = Migrate(app, db)
 
 
 @app.shell_context_processor
 def make_shell_context():
-    return dict(db=db, User=User, Role=Role)
+    return dict(db=db, User=User, Role=Role, Permission=Permission)
 
 
 @app.cli.command()

migrations/versions/968a609f1d46_.py

@@ -0,0 +1,32 @@
+"""empty message
+
+Revision ID: 968a609f1d46
+Revises: 
+Create Date: 2018-08-23 21:40:35.599873
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '968a609f1d46'
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('roles', sa.Column('default', sa.Boolean(), nullable=True))
+    op.add_column('roles', sa.Column('permissions', sa.Integer(), nullable=True))
+    op.create_index(op.f('ix_roles_default'), 'roles', ['default'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_roles_default'), table_name='roles')
+    op.drop_column('roles', 'permissions')
+    op.drop_column('roles', 'default')
+    # ### end Alembic commands ###

tests/test_user_model.py

@@ -1,5 +1,5 @@
 import unittest
-from app.models import User
+from app.models import User, Permission, AnonymousUser
 
 
 class UserModelTestCase(unittest.TestCase):
@@ -20,4 +20,20 @@ def test_password_verification(self):
     def test_password_salts_are_random(self):
         u = User(password='cat')
         u2 = User(password='cat')
-        self.assertTrue(u.password_hash != u2.password_hash)
+        self.assertTrue(u.password_hash != u2.password_hash)
+
+    def test_user_role(self):
+        u = User(email='[email protected]', password='woainvren1')
+        self.assertTrue(u.can(Permission.FOLLOW))
+        self.assertTrue(u.can(Permission.COMMENT))
+        self.assertTrue(u.can(Permission.WRITE))
+        self.assertTrue(u.can(Permission.MODERATE))
+        self.assertTrue(u.can(Permission.ADMIN))
+
+    def test_anonymous_user(self):
+        u = AnonymousUser()
+        self.assertFalse(u.can(Permission.FOLLOW))
+        self.assertFalse(u.can(Permission.COMMENT))
+        self.assertFalse(u.can(Permission.WRITE))
+        self.assertFalse(u.can(Permission.MODERATE))
+        self.assertFalse(u.can(Permission.ADMIN))