From 2908f96c883be75b5329001f46eeef3d47728bbd Mon Sep 17 00:00:00 2001
From: Eshaan Bansal <eshaan7bansal@gmail.com>
Date: Thu, 30 Apr 2020 21:25:29 +0530
Subject: [PATCH] admin views init moved into utils/

---
 src/FlaskRTBCTF/__init__.py            | 29 ++++-----------------
 src/FlaskRTBCTF/utils/admin_manager.py | 36 ++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/src/FlaskRTBCTF/__init__.py b/src/FlaskRTBCTF/__init__.py
index 20e3cff..413a938 100644
--- a/src/FlaskRTBCTF/__init__.py
+++ b/src/FlaskRTBCTF/__init__.py
@@ -3,12 +3,6 @@
 from flask import Flask
 
 from FlaskRTBCTF.config import Config
-from FlaskRTBCTF.admin.views import (
-    BaseModelView,
-    UserAdminView,
-    MachineAdminView,
-    NotificationAdminView,
-)
 from FlaskRTBCTF.utils import (
     db,
     bcrypt,
@@ -18,36 +12,23 @@
     mail,
     inject_app_context,
 )
-
-from FlaskRTBCTF.users.models import User, Logs
-from FlaskRTBCTF.main.models import Notification
-from FlaskRTBCTF.ctf.models import Machine
-
 from FlaskRTBCTF.users.routes import users
 from FlaskRTBCTF.ctf.routes import ctf
 from FlaskRTBCTF.main.routes import main
 
+
 _blueprints = (users, ctf, main)
 
+_extensions = (db, bcrypt, cache, login_manager, admin_manager, mail)
+
 
 def create_app(config_class=Config):
     app = Flask(__name__)
     app.config.from_object(Config)
-
-    db.init_app(app)
-    bcrypt.init_app(app)
-    cache.init_app(app)
-    login_manager.init_app(app)
-    admin_manager.init_app(app)
-    mail.init_app(app)
-
     app.context_processor(inject_app_context)
 
-    # Add model views for admin control
-    admin_manager.add_view(UserAdminView(User, db.session))
-    admin_manager.add_view(MachineAdminView(Machine, db.session))
-    admin_manager.add_view(NotificationAdminView(Notification, db.session))
-    admin_manager.add_view(BaseModelView(Logs, db.session))
+    for _ext in _extensions:
+        _ext.init_app(app)
 
     for _bp in _blueprints:
         app.register_blueprint(_bp)
diff --git a/src/FlaskRTBCTF/utils/admin_manager.py b/src/FlaskRTBCTF/utils/admin_manager.py
index 9dc0c25..93f4ed0 100644
--- a/src/FlaskRTBCTF/utils/admin_manager.py
+++ b/src/FlaskRTBCTF/utils/admin_manager.py
@@ -1,7 +1,43 @@
 from flask_admin import Admin
 from flask_admin.menu import MenuLink
 
+from FlaskRTBCTF.admin.views import (
+    BaseModelView,
+    UserAdminView,
+    MachineAdminView,
+    NotificationAdminView,
+    ChallengeAdminView,
+    UserChallengeAdminView,
+    UserMachineAdminView,
+)
+
+from FlaskRTBCTF.users.models import User, Logs
+from FlaskRTBCTF.main.models import Notification
+from FlaskRTBCTF.ctf.models import (
+    Machine,
+    Challenge,
+    Tag,
+    Category,
+    UserChallenge,
+    UserMachine,
+)
+from .models import db
 
 admin_manager = Admin(template_mode="bootstrap3")
 admin_manager.add_link(MenuLink(name="CTF Setup", url="/setup"))
 admin_manager.add_link(MenuLink(name="Go Back", url="/"))
+
+# Add model views for admin control
+admin_manager.add_view(NotificationAdminView(Notification, db.session))
+admin_manager.add_view(UserAdminView(User, db.session, category="Users"))
+admin_manager.add_view(BaseModelView(Logs, db.session, category="Users"))
+admin_manager.add_view(ChallengeAdminView(Challenge, db.session, category="Challenges"))
+admin_manager.add_view(
+    UserChallengeAdminView(UserChallenge, db.session, category="Challenges")
+)
+admin_manager.add_view(BaseModelView(Category, db.session, category="Challenges"))
+admin_manager.add_view(BaseModelView(Tag, db.session, category="Challenges"))
+admin_manager.add_view(MachineAdminView(Machine, db.session, category="Machines"))
+admin_manager.add_view(
+    UserMachineAdminView(UserMachine, db.session, category="Machines")
+)