AdvLabDB/advlabdb/customClasses.py

from flask import flash, redirect, request, url_for
from flask_admin import AdminIndexView, BaseView, expose
from flask_admin.contrib.sqla import ModelView
from flask_admin.model.helpers import get_mdict_item_or_list
from flask_admin.model.template import EndpointLinkRowAction
from flask_security import current_user
from sqlalchemy import and_

from advlabdb.dependent_funs import reportBadAttempt, userActiveSemester
from advlabdb.exceptions import DataBaseException, ModelViewException
from advlabdb.models import (
    Assistant,
    ExperimentMark,
    GroupExperiment,
    Part,
    PartStudent,
    SemesterExperiment,
)


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()

    @expose("/")
    def index(self):
        active_semester_experiment_marks_query = ExperimentMark.query.filter(
            ExperimentMark.part_student.has(PartStudent.part.has(Part.semester == userActiveSemester()))
        )

        number_of_all_experiment_marks = active_semester_experiment_marks_query.count()

        number_of_missing_final_experiment_marks = active_semester_experiment_marks_query.filter(
            ExperimentMark.final_experiment_mark == None
        ).count()

        return self.render(
            "admin_index.html",
            number_of_missing_final_experiment_marks=number_of_missing_final_experiment_marks,
            number_of_all_experiment_marks=number_of_all_experiment_marks,
        )


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

    @expose("/")
    def index(self):
        active_semester_experiment_marks_query = ExperimentMark.query.filter(
            ExperimentMark.group_experiment.has(
                GroupExperiment.semester_experiment.has(
                    and_(
                        SemesterExperiment.semester == userActiveSemester(),
                        SemesterExperiment.assistants.any(Assistant.user == current_user),
                    )
                )
            )
        )

        number_of_all_experiment_marks = active_semester_experiment_marks_query.count()

        number_of_missing_final_experiment_marks = active_semester_experiment_marks_query.filter(
            ExperimentMark.final_experiment_mark == None
        ).count()

        return self.render(
            "assistant_index.html",
            number_of_missing_final_experiment_marks=number_of_missing_final_experiment_marks,
            number_of_all_experiment_marks=number_of_all_experiment_marks,
        )


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

    can_view_details = False

    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"
    details_template = "admin_details.html"

    def is_accessible(self):
        return adminViewIsAccessible()


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

    can_edit = False
    column_display_actions = False

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

    """
    Every variable and method defined below in this class except queryFilter is NOT ALLOWED TO BE (completely) OVERWRITTEN!
    You can only extend the methods.
    queryFilter has to be implemented by overriding it.
    This is because of security reasons!
    """

    # Assistants are not allowed to create or delete.
    can_create = False
    can_delete = False

    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 NotImplementedError()

    def on_model_change(self, form, model, is_created):
        """
        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!")

    @expose("/edit/", methods=("GET", "POST"))
    def edit_view(self):
        """
        Prevent an assistant from seeing the edit form of a model not in his filter by using the GET method and entering an id.
        This is important if can_edit is set to True.
        """
        id = get_mdict_item_or_list(request.args, "id")

        if id is not None:
            model = self.get_one(id)

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

        return super().edit_view()

    @expose("/details/")
    def details_view(self):
        """
        Prevent an assistant from seeing the details of a model not in his filter by using the GET method and entering an id.
        This is important if can_view_details is set to True.
        """
        id = get_mdict_item_or_list(request.args, "id")

        if id is not None:
            model = self.get_one(id)

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

        return super().details_view()

    def get_details_columns(self):
        """
        Prevent showing unintended data if can_view_details is set to True and column_details_list was not set (which equals None).
        """
        if not self.column_details_list:
            self.column_details_list = self.column_list

        return super().get_details_columns()


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


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


class CustomIdEndpointLinkRowAction(EndpointLinkRowAction):
    def customId(self, row):
        raise NotImplementedError()

    def render(self, context, row_id, row):
        return super().render(context, self.customId(row), row)