You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

267 lines
9.1 KiB

  1. from datetime import datetime
  2. import hashlib
  3. from werkzeug.security import generate_password_hash, check_password_hash
  4. from itsdangerous import TimedJSONWebSignatureSerializer as Serializer
  5. from markdown import markdown
  6. import bleach
  7. from itsdangerous import BadSignature
  8. from flask import current_app
  9. from flask_login import UserMixin, AnonymousUserMixin
  10. from . import db, login_manager
  11. class Permission:
  12. FOLLOW = 1
  13. COMMENT = 2
  14. WRITE = 4
  15. MODERATE = 8
  16. ADMIN = 16
  17. class Role(db.Model):
  18. __tablename__ = 'roles'
  19. id = db.Column(db.Integer, primary_key=True)
  20. name = db.Column(db.String(64), unique=True)
  21. default = db.Column(db.Boolean, default=False, index=True)
  22. permissions = db.Column(db.Integer)
  23. users = db.relationship('User', backref='role', lazy='dynamic')
  24. def __init__(self, **kwargs):
  25. super(Role, self).__init__(**kwargs)
  26. if self.permissions is None:
  27. self.permissions = 0
  28. @staticmethod
  29. def insert_roles():
  30. roles = {
  31. 'User': [Permission.FOLLOW, Permission.COMMENT, Permission.WRITE],
  32. 'Moderator': [Permission.FOLLOW, Permission.COMMENT,
  33. Permission.WRITE, Permission.MODERATE],
  34. 'Administrator': [Permission.FOLLOW, Permission.COMMENT,
  35. Permission.WRITE, Permission.MODERATE,
  36. Permission.ADMIN]
  37. }
  38. default_role = 'User'
  39. for r in roles:
  40. role = Role.query.filter_by(name=r).first()
  41. if role is None:
  42. role = Role(name=r)
  43. role.reset_permissions()
  44. for perm in roles[r]:
  45. role.add_permission(perm)
  46. role.default = (role.name == default_role)
  47. db.session.add(role)
  48. db.session.commit()
  49. def add_permission(self, perm):
  50. if not self.has_permission(perm):
  51. self.permissions += perm
  52. def remove_permission(self, perm):
  53. if self.has_permission(perm):
  54. self.permissions -= perm
  55. def reset_permissions(self):
  56. self.permissions = 0
  57. def has_permission(self, perm):
  58. return self.permissions & perm == perm
  59. def __repr__(self):
  60. return '<Role %r>' % self.name
  61. class Follow(db.Model):
  62. __tablename__ = 'follows'
  63. follower_id = db.Column(db.Integer, db.ForeignKey('users.id'),
  64. primary_key=True)
  65. followed_id = db.Column(db.Integer, db.ForeignKey('users.id'),
  66. primary_key=True)
  67. timestamp = db.Column(db.DateTime, default=datetime.utcnow)
  68. class User(UserMixin, db.Model):
  69. __tablename__ = 'users'
  70. id = db.Column(db.Integer, primary_key=True)
  71. email = db.Column(db.String(64), unique=True, index=True)
  72. username = db.Column(db.String(64), unique=True, index=True)
  73. role_id = db.Column(db.Integer, db.ForeignKey('roles.id'))
  74. password_hash = db.Column(db.String(128))
  75. confirmed = db.Column(db.Boolean, default=False)
  76. name = db.Column(db.String(64))
  77. location = db.Column(db.String(64))
  78. about_me = db.Column(db.Text())
  79. member_since = db.Column(db.DateTime(), default=datetime.utcnow)
  80. last_seen = db.Column(db.DateTime(), default=datetime.utcnow)
  81. avatar_hash = db.Column(db.String(32))
  82. posts = db.relationship('Post', backref='author', lazy='dynamic')
  83. followed = db.relationship('Follow',
  84. foreign_keys=[Follow.follower_id],
  85. backref=db.backref('follower', lazy='joined'),
  86. lazy='dynamic',
  87. cascade='all, delete-orphan')
  88. followers = db.relationship('Follow',
  89. foreign_keys=[Follow.followed_id],
  90. backref=db.backref('followed', lazy='joined'),
  91. lazy='dynamic',
  92. cascade='all, delete-orphan')
  93. def __init__(self, **kwargs):
  94. super(User, self).__init__(**kwargs)
  95. if self.role is None:
  96. if self.email == current_app.config['FLASKY_ADMIN']:
  97. self.role = Role.query.filter_by(name='Administrator').first()
  98. else:
  99. self.role = Role.query.filter_by(default=True).first()
  100. if self.email is not None and self.avatar_hash is None:
  101. self.avatar_hash = self.gravatar_hash()
  102. @property
  103. def password(self):
  104. raise AttributeError('Password is not a readable attribute')
  105. @password.setter
  106. def password(self, password):
  107. self.password_hash = generate_password_hash(password)
  108. def verify_password(self, password):
  109. return check_password_hash(self.password_hash, password)
  110. def generate_confirmation_token(self, expiration=3600):
  111. s = Serializer(current_app.config['SECRET_KEY'], expiration)
  112. return s.dumps({'confirm': self.id}).decode('utf-8')
  113. def confirm(self, token):
  114. s = Serializer(current_app.config['SECRET_KEY'])
  115. try:
  116. data = s.loads(token.encode('utf-8'))
  117. except BadSignature:
  118. return False
  119. if data.get('confirm') != self.id:
  120. return False
  121. self.confirmed = True
  122. db.session.add(self)
  123. return True
  124. def generate_reset_token(self, expiration=3600):
  125. s = Serializer(current_app.config['SECRET_KEY'], expiration)
  126. return s.dumps({'reset': self.id}).decode('utf-8')
  127. @staticmethod
  128. def reset_password(token, new_password):
  129. s = Serializer(current_app.config['SECRET_KEY'])
  130. try:
  131. data = s.loads(token.encode('utf-8'))
  132. except BadSignature:
  133. return False
  134. user = User.query.get(data.get('reset'))
  135. if user is None:
  136. return False
  137. user.password = new_password
  138. db.session.add(user)
  139. return True
  140. def generate_email_change_token(self, new_email, expiration=3600):
  141. s = Serializer(current_app.config['SECRET_KEY'], expiration)
  142. return s.dumps(
  143. {'change_email': self.id, 'new_email': new_email}).decode('utf-8')
  144. def change_email(self, token):
  145. s = Serializer(current_app.config['SECRET_KEY'])
  146. try:
  147. data = s.loads(token.encode('utf-8'))
  148. except BadSignature:
  149. return False
  150. if data.get('change_email') != self.id:
  151. return False
  152. new_email = data.get('new_email')
  153. if new_email is None:
  154. return False
  155. if self.query.filter_by(email=new_email).first() is not None:
  156. return False
  157. self.email = new_email
  158. self.gravatar_hash = self.gravatar_hash()
  159. db.session.add(self)
  160. return True
  161. def can(self, perm):
  162. return self.role is not None and self.role.has_permission(perm)
  163. def is_administrator(self):
  164. return self.can(Permission.ADMIN)
  165. def ping(self):
  166. self.last_seen = datetime.utcnow()
  167. db.session.add(self)
  168. db.session.commit()
  169. def gravatar_hash(self):
  170. return hashlib.md5(self.email.lower().encode('utf-8')).hexdigest()
  171. def gravatar(self, size=100, default='identicon', rating='g'):
  172. url = 'https://secure.gravatar.com/avatar'
  173. hash = self.avatar_hash or self.gravatar_hash()
  174. return '{url}/{hash}?s={size}&d={default}&r={rating}'.format(
  175. url=url, hash=hash, size=size, default=default, rating=rating)
  176. def follow(self, user):
  177. if not self.is_following(user):
  178. f = Follow(follower=self, followed=user)
  179. db.session.add(f)
  180. def unfollow(self, user):
  181. f = self.followed.filter_by(followed_id=user.id).first()
  182. if f:
  183. db.session.delete(f)
  184. def is_following(self, user):
  185. if user.id is None:
  186. return False
  187. return self.followed.filter_by(
  188. followed_id=user.id).first() is not None
  189. def is_followed_by(self, user):
  190. if user.id is None:
  191. return False
  192. return self.followers.filter_by(
  193. follower_id=user.id).first() is not None
  194. def __repr__(self):
  195. return '<User %r>' % self.username
  196. class Post(db.Model):
  197. __tablename__ = 'posts'
  198. id = db.Column(db.Integer, primary_key=True)
  199. body = db.Column(db.Text)
  200. body_html = db.Column(db.Text)
  201. timestamp = db.Column(db.DateTime, index=True, default=datetime.utcnow)
  202. author_id = db.Column(db.Integer, db.ForeignKey('users.id'))
  203. @staticmethod
  204. def on_changed_body(target, value, oldvalue, initiator):
  205. allowed_tags = ['a', 'abbr', 'acronym', 'b', 'blockquote', 'code',
  206. 'em', 'i', 'li', 'ol', 'pre', 'strong', 'ul',
  207. 'h1', 'h2', 'h3', 'p']
  208. md = markdown(value, output_format='html')
  209. clean_md = bleach.clean(md, tags=allowed_tags, strip=True)
  210. target.body_html = bleach.linkify(clean_md)
  211. db.event.listen(Post.body, 'set', Post.on_changed_body)
  212. class AnonymousUser(AnonymousUserMixin):
  213. def can(self, perm):
  214. return False
  215. def is_administrator(self):
  216. return False
  217. login_manager.anonymous_user = AnonymousUser
  218. @login_manager.user_loader
  219. def load_user(user_id):
  220. return User.query.get(int(user_id))