class Role(db.Model):
__tablename__ = "roles"
id = db.Column(
db.Integer,
primary_key=True,
)
name = db.Column(
db.String(length=64),
unique=True,
)
users = db.relationship("User", back_populates="role", lazy="dynamic")
permissions = db.Column(db.Integer, )
def __init__(self, *args, **kwargs):
super(Role, self).__init__(*args, **kwargs)
if self.permissions is None:
self.permissions = 0b00000
def __repr__(self):
return "<Role {0}>".format(self.name)
def add_permission(self, permission):
if not self.has_permission(permission):
self.permissions += permission
def has_permission(self, permission):
return self.permissions & permission == permission
def reset_permission(self, ):
self.permissions = 0b00000
def remove_permission(self, permission):
if self.has_permission(permission):
self.permissions -= permission
class Image(db.Model):
__tablename__ = "images"
id = db.Column(
db.Integer,
primary_key=True,
)
user_id = db.Column(db.Integer(), db.ForeignKey('users.id'))
user = db.relationship("User", back_populates="photo")
class ShortUrl(db.Model):
__tablename__ = "short_urls"
id = db.Column(
db.Integer,
primary_key=True,
)
name = db.Column(
db.String(length=250),
unique=True,
)
url = db.Column(db.String(length=250), )
user_id = db.Column(db.Integer(), db.ForeignKey('users.id'))
user = db.relationship("User", back_populates="short_urls")
@staticmethod
def create_short_url(url, ):
return uuid.uuid4().hex
@staticmethod
def all_user_urls(user):
return ShortUrl.query.filter_by(user=user).all()
def __repr__(self):
return "<Short {0}>".format(self.name)
class User(UserMixin, db.Model):
__tablename__ = "users"
id = db.Column(
db.Integer,
primary_key=True,
)
name = db.Column(db.String(length=256, ), )
family = db.Column(db.String(length=256, ), )
username = db.Column(
db.String(length=256, ),
unique=True,
)
photo = db.relationship("Image", back_populates="user", uselist=False)
short_urls = db.relationship(
"ShortUrl",
back_populates="user",
)
password_hash = db.Column(db.String(length=256, ), )
role_id = db.Column(
db.Integer,
db.ForeignKey("roles.id"),
)
role = db.relationship(
"Role",
back_populates="users",
)
email = db.Column(db.String(length=256), unique=True)
location = db.Column(db.String(length=256))
about_me = db.Column(db.Text())
last_seen = db.Column(db.DateTime, default=datetime.utcnow)
is_active = db.Column(db.Boolean(), default=False)
def __init__(self, *args, **kwargs):
super(User, self).__init__(*args, **kwargs)
if self.role is None:
if self.email == current_app.config["ADMIN"]:
self.role = Role.query.filter_by(name="admin").first()
self.is_active = True
else:
self.role = Role.query.filter_by(name="user").first()
@property
def password(self):
raise AttributeError("[-] Setting password directly is not allowed!")
@password.setter
def password(self, password):
self.password_hash = generate_password_hash(password=password)
def verify_password(self, password):
return check_password_hash(self.password_hash, password)
def __repr__(self):
return "<User {0}>".format(self.name)
def can(self, permission):
if permission is not None:
return self.role.has_permission(permission)
return False
def update_last_seen(self, ):
self.last_seen = datetime.utcnow()
db.session.add(self)
db.session.commit()
def get_photo_url(self, ):
if self.photo:
photo = self.photo
return str(photo.id) + current_app.config.get("SAVE_EXTENSION")
def generate_signature(self, exp=60000):
serializer = TimedJSONWebSignatureSerializer(
secret_key=current_app.config['SECRET_KEY'],
salt='pd',
expires_in=exp,
)
return serializer.dumps({
'code': self.id,
})
def validate_signature(self, token):
result = False
serializer = TimedJSONWebSignatureSerializer(
secret_key=current_app.config["SECRET_KEY"],
salt='pd',
)
try:
data = serializer.loads(token)
if data.get('code') == self.id:
result = True
except Exception as e:
print(e)
return result
class User(UserMixin, db.Model):
id = db.Column(
db.Integer,
primary_key=True,
)
username = db.Column(
db.String(128),
unique=True,
)
password = db.Column(db.String(512), )
email = db.Column(
db.String(128),
unique=True,
)
posts = db.relationship("Post", backref="author", lazy="dynamic")
about_me = db.Column(db.String(256))
last_seen = db.Column(db.DateTime, default=datetime.utcnow)
followed = db.relationship(
"User",
secondary=followers,
primaryjoin=id == followers.c.follower,
secondaryjoin=id == followers.c.followed,
backref=db.backref("followers", lazy="dynamic"),
lazy="dynamic",
)
def set_password(self, password):
self.password = generate_password_hash(password, )
def check_password(self, password):
return check_password_hash(self.password, password)
def avatar(self, size):
email_hash = md5(self.email.lower().encode("utf-8")).hexdigest()
url = "https://www.gravatar.com/avatar/{}?d=robohash&s={}".format(
email_hash, size)
return url
def follow(self, user):
if not self.is_following(user):
self.followed.append(user)
def unfollow(self, user):
if self.is_following(user):
self.followed.remove(user)
def is_following(self, user):
if self.followed.filter_by(id=user.id).first():
return True
return False
def is_followed(self, user):
if self.followers.filter_by(id=user.id).first():
return True
return False
def get_followed_posts(self, ):
query = Post.query.join(followers,
(followers.c.followed == Post.user_id)).filter(
followers.c.follower == self.id).order_by(
Post.timestamp.desc())
return query
@staticmethod
def get_user_by_username(username, ):
user = User.query.filter_by(username=username).first()
return user
@staticmethod
def get_user_by_email(user_email, ):
user = User.query.filter_by(email=user_email).first()
return user
def set_lastseen(self, ):
self.last_seen = datetime.utcnow()
def build_jwt_token(
self,
exp=600,
):
payload = {
"exp": time() + exp,
"id": self.id,
}
key = current_app.config["SECRET_KEY"]
token = encode(payload=payload, key=key, algorithm='HS256')
token = token.decode("utf-8")
return token
@staticmethod
def get_user_by_token(token, ):
try:
key = current_app.config["SECRET_KEY"]
obj = decode(jwt=token, key=key, algorithms=[
"HS256",
])
id = obj.get("id")
return User.query.get(id)
except Exception as e:
return None
@staticmethod
def verify_user_token(token, ):
time_token = TimedJSONWebSignatureSerializer(
secret_key=current_app.config["SECRET_KEY"], )
try:
data = time_token.loads(token)
except (BadTimeSignature, SignatureExpired):
return None
else:
user_id = data["id"]
user = User.query.get(user_id)
return user
def build_token(self, ):
token = TimedJSONWebSignatureSerializer(
secret_key=current_app.config["SECRET_KEY"], expires_in=60)
return token.dumps({"id": self.id})
@staticmethod
def delete_user_by_id(user_id, ):
user = User.query.get(user_id)
db.session.delete(user)
db.session.commit()
@staticmethod
def update_user_by_id(user_id, data):
user = User.query.get(user_id)
for key, value in data.items():
setattr(user, key, value)
db.session.commit()
return user
def __repr__(self, ):
return '<%s.%s: %s, object at %s>' % (self.__class__.__module__,
self.__class__.__name__,
self.username, hex(id(self)))