Advertisement
Guest User

Untitled

a guest
Feb 12th, 2017
155
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 4.99 KB | None | 0 0
  1. from werkzeug.security import generate_password_hash, check_password_hash
  2. from itsdangerous import TimedJSONWebSignatureSerializer as Serializer
  3. from flask import current_app
  4. from flask_login import UserMixin, AnonymousUserMixin
  5. from . import db, login_manager
  6.  
  7.  
  8. class Permission:
  9.     FOLLOW = 0x01
  10.     COMMENT = 0x02
  11.     WRITE_ARTICLES = 0x04
  12.     MODERATE_COMMENTS = 0x08
  13.     ADMINISTER = 0x80
  14.  
  15.  
  16. class Role(db.Model):
  17.     __tablename__ = 'roles'
  18.     id = db.Column(db.Integer, primary_key=True)
  19.     name = db.Column(db.String(64), unique=True)
  20.     default = db.Column(db.Boolean, default=False, index=True)
  21.     permissions = db.Column(db.Integer)
  22.     users = db.relationship('User', backref='role', lazy='dynamic')
  23.  
  24.     @staticmethod
  25.     def insert_roles():
  26.         roles = {
  27.             'User': (Permission.FOLLOW |
  28.                      Permission.COMMENT |
  29.                      Permission.WRITE_ARTICLES, True),
  30.             'Moderator': (Permission.FOLLOW |
  31.                           Permission.COMMENT |
  32.                           Permission.WRITE_ARTICLES |
  33.                           Permission.MODERATE_COMMENTS, False),
  34.             'Administrator': (0xff, False)
  35.         }
  36.         for r in roles:
  37.             role = Role.query.filter_by(name=r).first()
  38.             if role is None:
  39.                 role = Role(name=r)
  40.             role.permissions = roles[r][0]
  41.             role.default = roles[r][1]
  42.             db.session.add(role)
  43.         db.session.commit()
  44.  
  45.     def __repr__(self):
  46.         return '<Role %r>' % self.name
  47.  
  48.  
  49. class User(UserMixin, db.Model):
  50.     __tablename__ = 'users'
  51.     id = db.Column(db.Integer, primary_key=True)
  52.     email = db.Column(db.String(64), unique=True, index=True)
  53.     username = db.Column(db.String(64), unique=True, index=True)
  54.     role_id = db.Column(db.Integer, db.ForeignKey('roles.id'))
  55.     password_hash = db.Column(db.String(128))
  56.     confirmed = db.Column(db.Boolean, default=False)
  57.  
  58.     def __init__(self, **kwargs):
  59.         super(User, self).__init__(**kwargs)
  60.         if self.role is None:
  61.             if self.email == current_app.config['FLASKY_ADMIN']:
  62.                 self.role = Role.query.filter_by(permissions=0xff).first()
  63.             if self.role is None:
  64.                 self.role = Role.query.filter_by(default=True).first()
  65.  
  66.     @property
  67.     def password(self):
  68.         raise AttributeError('password is not a readable attribute')
  69.  
  70.     @password.setter
  71.     def password(self, password):
  72.         self.password_hash = generate_password_hash(password)
  73.  
  74.     def verify_password(self, password):
  75.         return check_password_hash(self.password_hash, password)
  76.  
  77.     def generate_confirmation_token(self, expiration=3600):
  78.         s = Serializer(current_app.config['SECRET_KEY'], expiration)
  79.         return s.dumps({'confirm': self.id})
  80.  
  81.     def confirm(self, token):
  82.         s = Serializer(current_app.config['SECRET_KEY'])
  83.         try:
  84.             data = s.loads(token)
  85.         except:
  86.             return False
  87.         if data.get('confirm') != self.id:
  88.             return False
  89.         self.confirmed = True
  90.         db.session.add(self)
  91.         return True
  92.  
  93.     def generate_reset_token(self, expiration=3600):
  94.         s = Serializer(current_app.config['SECRET_KEY'], expiration)
  95.         return s.dumps({'reset': self.id})
  96.  
  97.     def reset_password(self, token, new_password):
  98.         s = Serializer(current_app.config['SECRET_KEY'])
  99.         try:
  100.             data = s.loads(token)
  101.         except:
  102.             return False
  103.         if data.get('reset') != self.id:
  104.             return False
  105.         self.password = new_password
  106.         db.session.add(self)
  107.         return True
  108.  
  109.     def generate_email_change_token(self, new_email, expiration=3600):
  110.         s = Serializer(current_app.config['SECRET_KEY'], expiration)
  111.         return s.dumps({'change_email': self.id, 'new_email': new_email})
  112.  
  113.     def change_email(self, token):
  114.         s = Serializer(current_app.config['SECRET_KEY'])
  115.         try:
  116.             data = s.loads(token)
  117.         except:
  118.             return False
  119.         if data.get('change_email') != self.id:
  120.             return False
  121.         new_email = data.get('new_email')
  122.         if new_email is None:
  123.             return False
  124.         if self.query.filter_by(email=new_email).first() is not None:
  125.             return False
  126.         self.email = new_email
  127.         db.session.add(self)
  128.         return True
  129.  
  130.     def can(self, permissions):
  131.         return self.role is not None and \
  132.             (self.role.permissions & permissions) == permissions
  133.  
  134.     def is_administrator(self):
  135.         return self.can(Permission.ADMINISTER)
  136.  
  137.     def __repr__(self):
  138.         return '<User %r>' % self.username
  139.  
  140.  
  141. class AnonymousUser(AnonymousUserMixin):
  142.     def can(self, permissions):
  143.         return False
  144.  
  145.     def is_administrator(self):
  146.         return False
  147.  
  148. login_manager.anonymous_user = AnonymousUser
  149.  
  150.  
  151. @login_manager.user_loader
  152. def load_user(user_id):
  153.     return User.query.get(int(user_id))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement