From 364273f1f12b341283fcb426865b8ae16f1c505e Mon Sep 17 00:00:00 2001 From: Kevin Chung Date: Wed, 6 May 2020 12:46:51 -0400 Subject: [PATCH 1/6] Add cascading delete constraint to DynamicChallenge --- CTFd/plugins/dynamic_challenges/__init__.py | 4 ++- ...8807ea_add_cascading_delete_to_dynamic_.py | 29 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) create mode 100644 migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py diff --git a/CTFd/plugins/dynamic_challenges/__init__.py b/CTFd/plugins/dynamic_challenges/__init__.py index f5edceb7..b02af073 100644 --- a/CTFd/plugins/dynamic_challenges/__init__.py +++ b/CTFd/plugins/dynamic_challenges/__init__.py @@ -239,7 +239,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( + None, 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) diff --git a/migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py b/migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py new file mode 100644 index 00000000..aaf6514e --- /dev/null +++ b/migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py @@ -0,0 +1,29 @@ +"""Add cascading delete to dynamic challenges + +Revision ID: b37fb68807ea +Revises: 1093835a1051 +Create Date: 2020-05-06 12:21:39.373983 + +""" +from alembic import op + + +# revision identifiers, used by Alembic. +revision = 'b37fb68807ea' +down_revision = '1093835a1051' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint(None, 'dynamic_challenge', type_='foreignkey') + op.create_foreign_key(None, 'dynamic_challenge', 'challenges', ['id'], ['id'], ondelete='CASCADE') + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint(None, 'dynamic_challenge', type_='foreignkey') + op.create_foreign_key(None, 'dynamic_challenge', 'challenges', ['id'], ['id']) + # ### end Alembic commands ### From 930da02231957f4151d9c8bfe2f62c67f077cc6b Mon Sep 17 00:00:00 2001 From: Kevin Chung Date: Wed, 6 May 2020 16:43:56 -0400 Subject: [PATCH 2/6] Run formatter --- ...37fb68807ea_add_cascading_delete_to_dynamic_.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py b/migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py index aaf6514e..af9a3fcc 100644 --- a/migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py +++ b/migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py @@ -9,21 +9,23 @@ from alembic import op # revision identifiers, used by Alembic. -revision = 'b37fb68807ea' -down_revision = '1093835a1051' +revision = "b37fb68807ea" +down_revision = "1093835a1051" branch_labels = None depends_on = None def upgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.drop_constraint(None, 'dynamic_challenge', type_='foreignkey') - op.create_foreign_key(None, 'dynamic_challenge', 'challenges', ['id'], ['id'], ondelete='CASCADE') + op.drop_constraint(None, "dynamic_challenge", type_="foreignkey") + op.create_foreign_key( + None, "dynamic_challenge", "challenges", ["id"], ["id"], ondelete="CASCADE" + ) # ### end Alembic commands ### def downgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.drop_constraint(None, 'dynamic_challenge', type_='foreignkey') - op.create_foreign_key(None, 'dynamic_challenge', 'challenges', ['id'], ['id']) + op.drop_constraint(None, "dynamic_challenge", type_="foreignkey") + op.create_foreign_key(None, "dynamic_challenge", "challenges", ["id"], ["id"]) # ### end Alembic commands ### From b5fe079922d851dd2ebf923a473f0d72406c2aed Mon Sep 17 00:00:00 2001 From: Kevin Chung Date: Wed, 6 May 2020 22:23:22 -0400 Subject: [PATCH 3/6] Create a migrations system for plugins --- CTFd/plugins/__init__.py | 1 + CTFd/plugins/dynamic_challenges/__init__.py | 6 +- ...8807ea_add_cascading_delete_to_dynamic_.py | 8 +-- CTFd/plugins/migrations.py | 65 +++++++++++++++++++ 4 files changed, 73 insertions(+), 7 deletions(-) rename {migrations/versions => CTFd/plugins/dynamic_challenges/migrations}/b37fb68807ea_add_cascading_delete_to_dynamic_.py (82%) create mode 100644 CTFd/plugins/migrations.py 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 b02af073..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 @@ -240,7 +241,7 @@ class DynamicValueChallenge(BaseChallenge): class DynamicChallenge(Challenges): __mapper_args__ = {"polymorphic_identity": "dynamic"} id = db.Column( - None, db.ForeignKey("challenges.id", ondelete="CASCADE"), primary_key=True + 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) @@ -252,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/migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py b/CTFd/plugins/dynamic_challenges/migrations/b37fb68807ea_add_cascading_delete_to_dynamic_.py similarity index 82% rename from migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py rename to CTFd/plugins/dynamic_challenges/migrations/b37fb68807ea_add_cascading_delete_to_dynamic_.py index af9a3fcc..2db3f9a5 100644 --- a/migrations/versions/b37fb68807ea_add_cascading_delete_to_dynamic_.py +++ b/CTFd/plugins/dynamic_challenges/migrations/b37fb68807ea_add_cascading_delete_to_dynamic_.py @@ -10,21 +10,21 @@ from alembic import op # revision identifiers, used by Alembic. revision = "b37fb68807ea" -down_revision = "1093835a1051" +down_revision = None branch_labels = None depends_on = None -def upgrade(): +def upgrade(op=None): # ### commands auto generated by Alembic - please adjust! ### - op.drop_constraint(None, "dynamic_challenge", type_="foreignkey") + op.drop_constraint("dynamic_challenge_ibfk_1", "dynamic_challenge", type_="foreignkey") op.create_foreign_key( None, "dynamic_challenge", "challenges", ["id"], ["id"], ondelete="CASCADE" ) # ### end Alembic commands ### -def downgrade(): +def downgrade(op=None): # ### commands auto generated by Alembic - please adjust! ### op.drop_constraint(None, "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..5393da35 --- /dev/null +++ b/CTFd/plugins/migrations.py @@ -0,0 +1,65 @@ +import os + +from alembic.migration import MigrationContext +from flask import current_app +from flask_migrate import Migrate, stamp +from sqlalchemy import create_engine +from alembic.operations import Operations + +from alembic.script import ScriptDirectory +from alembic.config import Config +from sqlalchemy import pool + +from CTFd.utils import get_config, set_config + +import inspect +import os + + +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 + caller_path = inspect.stack()[1].filename + 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) From c04235a5d11fd7704b6210334be7b0ee567df272 Mon Sep 17 00:00:00 2001 From: Kevin Chung Date: Thu, 7 May 2020 10:53:45 -0400 Subject: [PATCH 4/6] Fix some code and lints --- ...8807ea_add_cascading_delete_to_dynamic_.py | 9 +++--- CTFd/plugins/migrations.py | 32 ++++++++----------- 2 files changed, 18 insertions(+), 23 deletions(-) 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 index 2db3f9a5..6b4a35fc 100644 --- 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 @@ -1,13 +1,10 @@ """Add cascading delete to dynamic challenges Revision ID: b37fb68807ea -Revises: 1093835a1051 +Revises: Create Date: 2020-05-06 12:21:39.373983 """ -from alembic import op - - # revision identifiers, used by Alembic. revision = "b37fb68807ea" down_revision = None @@ -17,7 +14,9 @@ depends_on = None def upgrade(op=None): # ### commands auto generated by Alembic - please adjust! ### - op.drop_constraint("dynamic_challenge_ibfk_1", "dynamic_challenge", type_="foreignkey") + op.drop_constraint( + "dynamic_challenge_ibfk_1", "dynamic_challenge", type_="foreignkey" + ) op.create_foreign_key( None, "dynamic_challenge", "challenges", ["id"], ["id"], ondelete="CASCADE" ) diff --git a/CTFd/plugins/migrations.py b/CTFd/plugins/migrations.py index 5393da35..1cf28f65 100644 --- a/CTFd/plugins/migrations.py +++ b/CTFd/plugins/migrations.py @@ -1,20 +1,15 @@ -import os - -from alembic.migration import MigrationContext -from flask import current_app -from flask_migrate import Migrate, stamp -from sqlalchemy import create_engine -from alembic.operations import Operations - -from alembic.script import ScriptDirectory -from alembic.config import Config -from sqlalchemy import pool - -from CTFd.utils import get_config, set_config - 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") @@ -24,12 +19,13 @@ def upgrade(plugin_name=None): if plugin_name is None: # Get the directory name of the plugin if unspecified - caller_path = inspect.stack()[1].filename + # 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 - ) + engine = create_engine(database_url, poolclass=pool.NullPool) conn = engine.connect() context = MigrationContext.configure(conn) op = Operations(context) From 52d0c2719a0fc1b8d7be7d735cfbdfdce5da2056 Mon Sep 17 00:00:00 2001 From: Kevin Chung Date: Thu, 7 May 2020 11:23:48 -0400 Subject: [PATCH 5/6] Update migration --- ...8807ea_add_cascading_delete_to_dynamic_.py | 29 ++++++++++++++----- 1 file changed, 22 insertions(+), 7 deletions(-) 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 index 6b4a35fc..b28d5708 100644 --- 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 @@ -13,10 +13,17 @@ depends_on = None def upgrade(op=None): - # ### commands auto generated by Alembic - please adjust! ### - op.drop_constraint( - "dynamic_challenge_ibfk_1", "dynamic_challenge", type_="foreignkey" - ) + 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" ) @@ -24,7 +31,15 @@ def upgrade(op=None): def downgrade(op=None): - # ### commands auto generated by Alembic - please adjust! ### - op.drop_constraint(None, "dynamic_challenge", type_="foreignkey") + 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"]) - # ### end Alembic commands ### From 6f0c0b1a5282b1dc104fbb5ea863aebdd256257c Mon Sep 17 00:00:00 2001 From: Kevin Chung Date: Thu, 7 May 2020 12:27:04 -0400 Subject: [PATCH 6/6] Add test for dynamic challenge reset --- tests/challenges/test_dynamic.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) 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)