from flask_admin import AdminIndexView from flask_admin.contrib.sqla import ModelView from flask_security import current_user from flask import redirect, request, url_for, flash from advlabdb.exceptions import ModelViewException, DataBaseException def adminViewIsAccessible(): return current_user.has_role("admin") class SecureAdminIndexView(AdminIndexView): def is_accessible(self): return adminViewIsAccessible() 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 SecureModelView(ModelView): can_export = True can_set_page_size = True create_modal = True edit_modal = True details_modal = True list_template = "admin_list.html" create_template = "admin_create.html" edit_template = "admin_edit.html" queryFilter = None customCreateModel = None def is_accessible(self): return adminViewIsAccessible() 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