models.py 2.11 KB
from pprint import pprint

from werkzeug.security import generate_password_hash, check_password_hash
from flask_login import UserMixin
from app.models import db

#
# Roles
#

ADMIN = 50
PROJECT = 40
SERVICE = 30
AGENT = 10
PUBLIC = 0

_roleToName = {
    ADMIN: 'ADMIN',
    PROJECT: 'PROJECT',
    SERVICE: 'SERVICE',
    AGENT: 'AGENT',
    PUBLIC: 'PUBLIC',
}
_nameToRole = {
    'ADMIN': ADMIN,
    'PROJECT': PROJECT,
    'SERVICE': SERVICE,
    'AGENT': AGENT,
    'PUBLIC': PUBLIC,
}


def _check_role(role):
    if isinstance(role, int):
        rv = role
    elif str(role) == role:
        role = role.upper()
        if role not in _nameToRole:
            raise ValueError("Unknown role: %r" % role)
        rv = _nameToRole[role]
    else:
        raise TypeError("Role not an integer or a valid string: %r" % role)
    return rv


class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)  # primary keys are required by SQLAlchemy
    email = db.Column(db.String(100), unique=True)
    name = db.Column(db.String(100))
    login = db.Column(db.String(100), unique=True)
    role = db.Column(db.Integer, default=0)
    password = db.Column(db.String(128))
    password_hash = db.Column(db.String(128))

    def __repr__(self):
        return "i: {}, n: {}, e: {}, l: {}".format(self.id, self.name, self.email, self.login)

    # Set role at construction time
    def __init__(self, **kwargs):
        super(User, self).__init__(**kwargs)
        if ('role' in kwargs):
            self.set_role(kwargs['role'])
        if ('password' in kwargs):
            self.set_password(kwargs['password'])
            self.password = None

    def set_role(self, role):
        self.role = _check_role(role)

    def has_role(self, role):
        role = _check_role(role)
        return self.role == role

    def has_role_or_higher(self, role):
        role = _check_role(role)
        return self.role and (self.role >= role)

    def set_password(self, password):
        self.password_hash = generate_password_hash(password)

    def check_password(self, password):
        return check_password_hash(self.password_hash, password)