diff --git a/CTFd/plugins/__init__.py b/CTFd/plugins/__init__.py index 42da0da5..7e97824e 100644 --- a/CTFd/plugins/__init__.py +++ b/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__) + "/*")) diff --git a/CTFd/plugins/dynamic_challenges/__init__.py b/CTFd/plugins/dynamic_challenges/__init__.py index f5edceb7..04ebc17f 100644 --- a/CTFd/plugins/dynamic_challenges/__init__.py +++ b/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/" diff --git a/CTFd/plugins/dynamic_challenges/migrations/b37fb68807ea_add_cascading_delete_to_dynamic_.py b/CTFd/plugins/dynamic_challenges/migrations/b37fb68807ea_add_cascading_delete_to_dynamic_.py new file mode 100644 index 00000000..b28d5708 --- /dev/null +++ b/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"]) diff --git a/CTFd/plugins/migrations.py b/CTFd/plugins/migrations.py new file mode 100644 index 00000000..1cf28f65 --- /dev/null +++ b/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) diff --git a/tests/challenges/test_dynamic.py b/tests/challenges/test_dynamic.py index e34bb0e4..04953520 100644 --- a/tests/challenges/test_dynamic.py +++ b/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)