Merge pull request #1388 from CTFd/1386-proper-deletion-constraint-for-dynamics

* Add cascading delete constraint to `DynamicChallenge` to help with Reset functionality
* Add a system for running migrations from within plugins
* Closes #1386

CTFd/plugins/__init__.py

@@ -176,6 +176,7 @@ def init_plugins(app):
 
     app.admin_plugin_menu_bar = []
     app.plugin_menu_bar = []
+    app.plugins_dir = os.path.dirname(__file__)
 
     if app.config.get("SAFE_MODE", False) is False:
         modules = sorted(glob.glob(os.path.dirname(__file__) + "/*"))

CTFd/plugins/dynamic_challenges/__init__.py

@@ -15,6 +15,7 @@ from CTFd.models import (
     db,
 )
 from CTFd.plugins import register_plugin_assets_directory
+from CTFd.plugins.migrations import upgrade
 from CTFd.plugins.challenges import CHALLENGE_CLASSES, BaseChallenge
 from CTFd.plugins.flags import get_flag_class
 from CTFd.utils.modes import get_model
@@ -239,7 +240,9 @@ class DynamicValueChallenge(BaseChallenge):
 
 class DynamicChallenge(Challenges):
     __mapper_args__ = {"polymorphic_identity": "dynamic"}
-    id = db.Column(None, db.ForeignKey("challenges.id"), primary_key=True)
+    id = db.Column(
+        db.Integer, db.ForeignKey("challenges.id", ondelete="CASCADE"), primary_key=True
+    )
     initial = db.Column(db.Integer, default=0)
     minimum = db.Column(db.Integer, default=0)
     decay = db.Column(db.Integer, default=0)
@@ -250,8 +253,7 @@ class DynamicChallenge(Challenges):
 
 
 def load(app):
-    # upgrade()
-    app.db.create_all()
+    upgrade()
     CHALLENGE_CLASSES["dynamic"] = DynamicValueChallenge
     register_plugin_assets_directory(
         app, base_path="/plugins/dynamic_challenges/assets/"

CTFd/plugins/dynamic_challenges/migrations/b37fb68807ea_add_cascading_delete_to_dynamic_.py

@@ -0,0 +1,45 @@
+"""Add cascading delete to dynamic challenges
+
+Revision ID: b37fb68807ea
+Revises:
+Create Date: 2020-05-06 12:21:39.373983
+
+"""
+# revision identifiers, used by Alembic.
+revision = "b37fb68807ea"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade(op=None):
+    bind = op.get_bind()
+    url = str(bind.engine.url)
+    if url.startswith("mysql"):
+        op.drop_constraint(
+            "dynamic_challenge_ibfk_1", "dynamic_challenge", type_="foreignkey"
+        )
+    elif url.startswith("postgres"):
+        op.drop_constraint(
+            "dynamic_challenge_id_fkey", "dynamic_challenge", type_="foreignkey"
+        )
+
+    op.create_foreign_key(
+        None, "dynamic_challenge", "challenges", ["id"], ["id"], ondelete="CASCADE"
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade(op=None):
+    bind = op.get_bind()
+    url = str(bind.engine.url)
+    if url.startswith("mysql"):
+        op.drop_constraint(
+            "dynamic_challenge_ibfk_1", "dynamic_challenge", type_="foreignkey"
+        )
+    elif url.startswith("postgres"):
+        op.drop_constraint(
+            "dynamic_challenge_id_fkey", "dynamic_challenge", type_="foreignkey"
+        )
+
+    op.create_foreign_key(None, "dynamic_challenge", "challenges", ["id"], ["id"])

CTFd/plugins/migrations.py

@@ -0,0 +1,61 @@
+import inspect
+import os
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.operations import Operations
+from alembic.script import ScriptDirectory
+from flask import current_app
+from sqlalchemy import create_engine, pool
+
+from CTFd.utils import get_config, set_config
+
+
+def upgrade(plugin_name=None):
+    database_url = current_app.config.get("SQLALCHEMY_DATABASE_URI")
+    if database_url.startswith("sqlite"):
+        current_app.db.create_all()
+        return
+
+    if plugin_name is None:
+        # Get the directory name of the plugin if unspecified
+        # Doing it this way doesn't waste the rest of the inspect.stack call
+        frame = inspect.currentframe()
+        caller_info = inspect.getframeinfo(frame.f_back)
+        caller_path = caller_info[0]
+        plugin_name = os.path.basename(os.path.dirname(caller_path))
+
+    engine = create_engine(database_url, poolclass=pool.NullPool)
+    conn = engine.connect()
+    context = MigrationContext.configure(conn)
+    op = Operations(context)
+
+    # Find the list of migrations to run
+    config = Config()
+    config.set_main_option(
+        "script_location",
+        os.path.join(current_app.plugins_dir, plugin_name, "migrations"),
+    )
+    config.set_main_option(
+        "version_locations",
+        os.path.join(current_app.plugins_dir, plugin_name, "migrations"),
+    )
+    script = ScriptDirectory.from_config(config)
+
+    # get current revision for plugin
+    lower = get_config(plugin_name + "_alembic_version")
+    upper = script.get_current_head()
+
+    # Apply from lower to upper
+    revs = list(script.iterate_revisions(lower=lower, upper=upper))
+    revs.reverse()
+
+    try:
+        for r in revs:
+            with context.begin_transaction():
+                r.module.upgrade(op=op)
+    finally:
+        conn.close()
+
+    # Set the new latest revision
+    set_config(plugin_name + "_alembic_version", upper)

tests/challenges/test_dynamic.py

@@ -47,7 +47,6 @@ def test_can_create_dynamic_challenge():
 
 
 def test_can_update_dynamic_challenge():
-    """Test that dynamic challenges can be deleted"""
     app = create_ctfd(enable_plugins=True)
     with app.app_context():
         challenge_data = {
@@ -150,6 +149,7 @@ def test_can_add_requirement_dynamic_challenge():
 
 
 def test_can_delete_dynamic_challenge():
+    """Test that dynamic challenges can be deleted"""
     app = create_ctfd(enable_plugins=True)
     with app.app_context():
         register_user(app)
@@ -317,3 +317,33 @@ def test_dynamic_challenge_value_isnt_affected_by_hidden_users():
                 chal = DynamicChallenge.query.filter_by(id=1).first()
                 assert chal.value == chal.initial
     destroy_ctfd(app)
+
+
+def test_dynamic_challenges_reset():
+    app = create_ctfd(enable_plugins=True)
+    with app.app_context():
+        client = login_as_user(app, name="admin", password="password")
+
+        challenge_data = {
+            "name": "name",
+            "category": "category",
+            "description": "description",
+            "value": 100,
+            "decay": 20,
+            "minimum": 1,
+            "state": "hidden",
+            "type": "dynamic",
+        }
+
+        r = client.post("/api/v1/challenges", json=challenge_data)
+        assert Challenges.query.count() == 1
+        assert DynamicChallenge.query.count() == 1
+
+        with client.session_transaction() as sess:
+            data = {"nonce": sess.get("nonce"), "challenges": "on"}
+            r = client.post("/admin/reset", data=data)
+            assert r.location.endswith("/admin/statistics")
+        assert Challenges.query.count() == 0
+        assert DynamicChallenge.query.count() == 0
+
+    destroy_ctfd(app)