implemented caching and other fixes

Author: eshaan7

app.json

@@ -5,6 +5,9 @@
   "addons": [
     {
       "plan": "heroku-postgresql"
+    },
+    {
+      "plan": "heroku-redis"
     }
   ],
   "buildpacks": [
@@ -21,6 +24,9 @@
       "description": "Administrator password",
       "generator": "secret"
     },
+    "ADMIN_EMAIL": {
+      "description": "Administrator email"
+    },
     "MAIL_USER": {
       "description": "Username for mail service",
       "required": false

src/FlaskRTBCTF/__init__.py

@@ -3,10 +3,16 @@
 from flask import Flask
 
 from FlaskRTBCTF.config import Config
-from FlaskRTBCTF.admin.views import BaseModelView, UserAdminView, MachineAdminView
+from FlaskRTBCTF.admin.views import (
+    BaseModelView,
+    UserAdminView,
+    MachineAdminView,
+    NotificationAdminView,
+)
 from FlaskRTBCTF.utils import (
     db,
     bcrypt,
+    cache,
     login_manager,
     admin_manager,
     mail,
@@ -30,6 +36,7 @@ def create_app(config_class=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)
@@ -39,7 +46,7 @@ def create_app(config_class=Config):
     # 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(BaseModelView(Notification, db.session))
+    admin_manager.add_view(NotificationAdminView(Notification, db.session))
     admin_manager.add_view(BaseModelView(Logs, db.session))
 
     for _bp in _blueprints:

src/FlaskRTBCTF/admin/views.py

@@ -1,6 +1,6 @@
 """ Admin Model Views. """
 
-from flask import abort, redirect, flash
+from flask import abort, redirect, flash, url_for, request
 from flask_login import current_user
 from flask_admin import expose
 from flask_admin.form import SecureForm
@@ -32,6 +32,7 @@ def _handle_view(self, name, **kwargs):
 
 class UserAdminView(BaseModelView):
     column_exclude_list = ("password",)
+    form_exclude_list = ("password",)
     column_searchable_list = ("username", "email")
 
     @expose("/new/")
@@ -42,6 +43,17 @@ def create_view(self):
 
 class MachineAdminView(BaseModelView):
     column_searchable_list = ("name", "ip")
-    form_choices = {
-        "hardness": [("easy", "Easy"), ("medium", "Medium"), ("hard", "Hard")]
-    }
+
+    @expose("/new/")
+    def create_view(self):
+        return redirect(url_for("ctf.new_machine"))
+
+    @expose("/edit/")
+    def edit_view(self):
+        id = int(request.args["id"])
+        return redirect(url_for("ctf.edit_machine", id=id))
+
+
+class NotificationAdminView(BaseModelView):
+    column_searchable_list = ("title",)
+    form_excluded_columns = ("timestamp",)

src/FlaskRTBCTF/config.py

@@ -7,14 +7,14 @@
 
 
 class Config:
+    DEBUG = False  # Turn DEBUG OFF before deployment
     SECRET_KEY = handle_secret_key()
     SQLALCHEMY_DATABASE_URI = os.environ.get("DATABASE_URL") or "sqlite:///site.db"
     # For local use, one can simply use SQLlite with: 'sqlite:///site.db'
     # For deployment on Heroku use: `os.environ.get('DATABASE_URL')`
     # in all other cases: `os.environ.get('SQLALCHEMY_DATABASE_URI')`
-    FLASK_ADMIN_SWATCH = ["journal", "paper", "yeti", "cosmo"][2]
     SQLALCHEMY_TRACK_MODIFICATIONS = False
-    DEBUG = False  # Turn DEBUG OFF before deployment
+    FLASK_ADMIN_SWATCH = ("journal", "paper", "yeti", "cosmo")[3]
     # TEMPLATES_AUTO_RELOAD = True
     MAIL_SERVER = "smtp.googlemail.com"
     MAIL_PORT = 587

src/FlaskRTBCTF/ctf/forms.py

@@ -1,9 +1,42 @@
 from flask_wtf import FlaskForm
-from wtforms import StringField, SubmitField, HiddenField
-from wtforms.validators import DataRequired, Length, ValidationError
+from wtforms import StringField, SubmitField, HiddenField, RadioField
+from wtforms.validators import DataRequired, Length, ValidationError, IPAddress
+from wtforms.fields.html5 import IntegerField
 from .models import Machine
 
 
+class MachineForm(FlaskForm):
+    name = StringField("Name", validators=[DataRequired(), Length(min=4, max=32)])
+    os = RadioField(
+        "Operating System of machine",
+        validators=[DataRequired()],
+        choices=(("linux", "Linux"), ("windows", "Windows"), ("android", "Android")),
+    )
+    user_hash = StringField(
+        "User Hash", validators=[DataRequired(), Length(min=32, max=32)]
+    )
+    root_hash = StringField(
+        "Root Hash", validators=[DataRequired(), Length(min=32, max=32)]
+    )
+    user_points = IntegerField("Points for User Hash", validators=[DataRequired()])
+    root_points = IntegerField("Points for Root Hash", validators=[DataRequired()])
+    ip = StringField(
+        "IPv4 address of machine", validators=[DataRequired(), IPAddress()]
+    )
+    hardness = RadioField(
+        "Difficuly Level",
+        validators=[DataRequired()],
+        choices=(
+            ("easy", "Easy"),
+            ("medium", "Medium"),
+            ("hard", "Hard"),
+            ("insane", "Insane"),
+        ),
+    )
+
+    submit = SubmitField("Submit")
+
+
 class UserHashForm(FlaskForm):
     machine_id = HiddenField("Machine ID", validators=[DataRequired()])
     user_hash = StringField(
@@ -30,7 +63,7 @@ def validate_root_hash(self, root_hash):
         box = Machine.query.get(int(self.machine_id.data))
         if not box:
             raise ValidationError("No machine with that ID exists")
-        elif box.user_hash == str(root_hash.data):
+        elif box.root_hash == str(root_hash.data):
             pass
         else:
             raise ValidationError("Incorrect Root Hash.")

src/FlaskRTBCTF/ctf/models.py

@@ -1,4 +1,4 @@
-from FlaskRTBCTF.utils import db
+from FlaskRTBCTF.utils import db, cache
 
 # Machine Table
 
@@ -11,6 +11,12 @@ class Machine(db.Model):
     root_hash = db.Column(db.String(32), nullable=False)
     user_points = db.Column(db.Integer, default=0)
     root_points = db.Column(db.Integer, default=0)
-    os = db.Column(db.String(16), nullable=False)
-    ip = db.Column(db.String(45), nullable=False)
-    hardness = db.Column(db.String(16), nullable=False, default="Easy")
+    os = db.Column(db.String, nullable=False, default="linux")
+    ip = db.Column(db.String(64), nullable=False)
+    hardness = db.Column(db.String, nullable=False, default="Easy")
+
+    @staticmethod
+    @cache.cached(timeout=3600, key_prefix="machines")
+    def get_all():
+        _machines = Machine.query.all()
+        return _machines

src/FlaskRTBCTF/ctf/routes.py

@@ -6,39 +6,28 @@
 from flask import Blueprint, render_template, flash, request, redirect, url_for
 from flask_login import current_user, login_required
 
-from FlaskRTBCTF import db
 from FlaskRTBCTF.users.models import User, Logs
-from FlaskRTBCTF.utils import is_past_running_time
+from FlaskRTBCTF.utils import db, cache, is_past_running_time, admin_only
 from .models import Machine
-from .forms import UserHashForm, RootHashForm
+from .forms import UserHashForm, RootHashForm, MachineForm
 
 
 ctf = Blueprint("ctf", __name__)
 
 
-# context processor
-
-
-@ctf.context_processor
-def inject_context():
-    boxes = Machine.query.all()
-    past_running_time = is_past_running_time()
-
-    return dict(boxes=boxes, past_running_time=past_running_time)
-
-
 # Scoreboard
 
 
 @ctf.route("/scoreboard")
-@login_required
+@cache.cached(timeout=120, key_prefix="scoreboard")
 def scoreboard():
-    users_score = User.query.order_by(User.points.desc()).all()
-    userNameScoreList = []
-    for u in users_score:
-        userNameScoreList.append({"username": u.username, "score": u.points})
+    users_scores = (
+        User.query.with_entities(User.username, User.points)
+        .order_by(User.points.desc())
+        .all()
+    )
 
-    return render_template("scoreboard.html", scores=userNameScoreList)
+    return render_template("scoreboard.html", scores=users_scores)
 
 
 # Machines Info
@@ -50,7 +39,11 @@ def machines():
     userHashForm = UserHashForm()
     rootHashForm = RootHashForm()
 
+    boxes = Machine.get_all()
+    past_running_time = is_past_running_time()
+
     if request.method == "GET":
+
         log = Logs.query.get(current_user.id)
 
         # check if it is the first visit to machine page for user
@@ -60,50 +53,106 @@ def machines():
             db.session.commit()
 
     else:
-        if is_past_running_time():
+        if past_running_time:
             flash("Sorry! CTF has ended.", "danger")
             return redirect(url_for("ctf.machines"))
 
         """
            Todo: Get Object from UserMachine Model, dummy object given below
         """
-        user_machine: object = {
-            "machine_id": 1,
-            "user_id": 1,
-            "owned_user": False,
-            "owned_root": False,
-        }
-
-        if user_machine.owned_user:
-            flash("You already own User.", "success")
-            return redirect(url_for("ctf.machines"))
-
-        elif user_machine.owned_root:
-            flash("You already own System.", "success")
-            return redirect(url_for("ctf.machines"))
-
-        elif userHashForm.submit_user_hash.data and userHashForm.validate_on_submit():
+        # user_machine: object = {
+        #     "machine_id": 1,
+        #     "user_id": 1,
+        #     "owned_user": False,
+        #     "owned_root": False,
+        # }
+
+        # if user_machine.owned_user:
+        #     flash("You already own User.", "success")
+        #     return redirect(url_for("ctf.machines"))
+
+        # elif user_machine.owned_root:
+        #     flash("You already own System.", "success")
+        #     return redirect(url_for("ctf.machines"))
+
+        if userHashForm.submit_user_hash.data and userHashForm.validate_on_submit():
             box = Machine.query.get(int(userHashForm.machine_id.data))
-            user_machine.owned_user = True
+            # user_machine.owned_user = True
             current_user.points += box.user_points
             log = Logs.query.get(current_user.id)
             log.userSubmissionIP = request.access_route[0]
             log.userSubmissionTime = datetime.utcnow()
             log.userOwnTime = str(log.userSubmissionTime - log.machineVisitTime)
             db.session.commit()
+            cache.delete(key="scoreboard")
             flash("Congrats! correct user hash.", "success")
 
         elif rootHashForm.submit_root_hash.data and rootHashForm.validate_on_submit():
             box = Machine.query.get(int(rootHashForm.machine_id.data))
-            user_machine.owned_root = True
+            # user_machine.owned_root = True
             current_user.points += box.root_points
             log = Logs.query.get(current_user.id)
             log.rootSubmissionIP = request.access_route[0]
             log.rootSubmissionTime = datetime.utcnow()
             log.rootOwnTime = str(log.rootSubmissionTime - log.machineVisitTime)
             db.session.commit()
+            cache.delete(key="scoreboard")
             flash("Congrats! correct root hash.", "success")
 
+        else:
+            errors = userHashForm.user_hash.errors or rootHashForm.root_hash.errors
+            for e in errors:
+                flash(e, "danger")
+
+        return redirect(url_for("ctf.machines"))
+
     return render_template(
-        "machine.html", userHashForm=userHashForm, rootHashForm=rootHashForm,
+        "machines.html",
+        boxes=boxes,
+        past_running_time=past_running_time,
+        userHashForm=userHashForm,
+        rootHashForm=rootHashForm,
     )
+
+
+@ctf.route("/machines/new", methods=["GET", "POST"])
+@admin_only
+def new_machine():
+    form = MachineForm(obj=Machine.query.get(1))
+    if request.method == "GET":
+        return render_template(
+            "new_machine.html", form_title="Add New Machine", form=form
+        )
+    else:
+        if form.validate_on_submit():
+            new_machine = Machine()
+            form.populate_obj(new_machine)
+            db.session.add(new_machine)
+            db.session.commit()
+            cache.delete(key="machines")
+            flash(f"{form.name.data} has been added.", "success")
+            return redirect(url_for("ctf.machines"))
+        else:
+            flash(form.errors, "danger")
+            return redirect(request.url)
+
+
+@ctf.route("/machines/edit/<int:id>", methods=["GET", "POST"])
+@admin_only
+def edit_machine(id):
+    machine = Machine.query.get_or_404(id)
+    form = MachineForm(obj=machine)
+    if request.method == "GET":
+        return render_template(
+            "new_machine.html", form_title=f"Editing machine #{id}", form=form
+        )
+    else:
+        if form.validate_on_submit():
+            form.populate_obj(machine)
+            db.session.commit()
+            cache.delete(key="machines")
+            flash(f"{form.name.data} has been edited.", "success")
+            return redirect(url_for("ctf.machines"))
+        else:
+            flash(form.errors, "danger")
+            return redirect(request.url)