AdvLabDB/advlabdb/customClasses.py

from flask import flash, redirect, request, url_for
from flask_admin import AdminIndexView, BaseView
from flask_admin.contrib.sqla import ModelView
from flask_security import current_user

from advlabdb.exceptions import DataBaseException, ModelViewException
from advlabdb.utils import reportBadAttempt


def adminViewIsAccessible():
    return current_user.has_role("admin")


def assistantViewIsAccessible():
    return current_user.has_role("assistant")


class CustomIndexView(AdminIndexView):
    def inaccessible_callback(self, name, **kwargs):
        # Redirect to login page if user doesn't have access
        return redirect(url_for("security.login", next=request.url))


class SecureAdminIndexView(CustomIndexView):
    def is_accessible(self):
        return adminViewIsAccessible()


class SecureAssistantIndexView(CustomIndexView):
    def is_accessible(self):
        return assistantViewIsAccessible()


class CustomModelView(ModelView):
    create_modal = True
    edit_modal = True
    details_modal = True

    queryFilter = None
    customCreateModel = None

    def inaccessible_callback(self, name, **kwargs):
        # Redirect to login page if user doesn't have access
        return redirect(url_for("security.login", next=request.url))

    def get_query(self):
        if self.queryFilter:
            return super().get_query().filter(self.queryFilter())
        else:
            return super().get_query()

    def get_count_query(self):
        if self.queryFilter:
            return super().get_count_query().filter(self.queryFilter())
        else:
            return super().get_count_query()

    def handle_view_exception(self, exc):
        if type(exc) in (ModelViewException, DataBaseException):
            flash(str(exc), "error")
            return True

        return super().handle_view_exception(exc)

    def create_model(self, form):
        if not self.customCreateModel:
            return super().create_model(form)
        else:
            try:
                model = self.customCreateModel(form)

                self.session.add(model)

                self.on_model_change(form, model, True)

                self.session.commit()
            except Exception as ex:
                flash(str(ex), "error")

                self.session.rollback()
            else:
                self.after_model_change(form, model, True)

                return model


class SecureAdminModelView(CustomModelView):
    can_export = True
    can_set_page_size = True

    can_create = True
    can_edit = True
    can_delete = True
    column_display_actions = True

    list_template = "admin_list.html"
    create_template = "admin_create.html"
    edit_template = "admin_edit.html"

    def is_accessible(self):
        return adminViewIsAccessible()


class SecureAssistantModelView(CustomModelView):
    can_export = False
    can_set_page_size = False

    can_create = False
    can_edit = False
    can_delete = False
    column_display_actions = False

    list_template = "assistant_list.html"
    create_template = "assistant_create.html"
    edit_template = "assistant_edit.html"

    def is_accessible(self):
        return assistantViewIsAccessible()

    def queryFilter(self):
        """
        A default filter has to be implemented to restrict assistants read/write access.
        See on_model_change!
        """
        raise ModelViewException("Not implemented!")

    def on_model_change(self, form, model, is_created):
        """
        This method is NOT ALLOWED TO BE (completely) OVERWRITTEN!
        This method uses the filter returned by queryFilter (which has to be implemented!) to prevent assistants
        from modifying models not listed on their view by sending a POST request with a different id.
        You can extend this method by implementing a custom on_model_change and then calling super().on_model_change within it.
        """
        if is_created:
            reportBadAttempt("An assistant tried to create a model!")
            raise ModelViewException("Assistants can not create models!")

        if model not in self.get_query():
            reportBadAttempt("An assistant tried to change a model not in his filter!")
            raise ModelViewException("Unauthorized action!")

    def on_model_delete(self, model):
        reportBadAttempt("An assistant tried to delete a model!")
        raise ModelViewException("Assistants can not delete models!")


class SecureAdminBaseView(BaseView):
    def is_accessible(self):
        return adminViewIsAccessible()


class SecureAssistantBaseView(BaseView):
    def is_accessible(self):
        return assistantViewIsAccessible()