Repository: CTFd/CTFd

Issue:
More granular reset
People often request the ability to reset individual tables instead of the current reset approach


Discussion:

Resolving patch:
diff --git a/CTFd/__init__.py b/CTFd/__init__.py
--- a/CTFd/__init__.py
+++ b/CTFd/__init__.py
@@ -179,6 +179,18 @@ def create_app(config="CTFd.config.Config"):
 
         # Alembic sqlite support is lacking so we should just create_all anyway
         if url.drivername.startswith("sqlite"):
+            # Enable foreign keys for SQLite. This must be before the
+            # db.create_all call because tests use the in-memory SQLite
+            # database (each connection, including db creation, is a new db).
+            from sqlalchemy.engine import Engine
+            from sqlalchemy import event
+
+            @event.listens_for(Engine, "connect")
+            def set_sqlite_pragma(dbapi_connection, connection_record):
+                cursor = dbapi_connection.cursor()
+                cursor.execute("PRAGMA foreign_keys=ON")
+                cursor.close()
+
             db.create_all()
             stamp_latest_revision()
         else:
diff --git a/CTFd/admin/__init__.py b/CTFd/admin/__init__.py
--- a/CTFd/admin/__init__.py
+++ b/CTFd/admin/__init__.py
@@ -14,10 +14,13 @@
     url_for,
 )
 
-from CTFd.cache import cache, clear_config
+from CTFd.cache import cache, clear_config, clear_standings, clear_pages
 from CTFd.models import (
     Awards,
+    Challenges,
     Configs,
+    Notifications,
+    Pages,
     Solves,
     Submissions,
     Teams,
@@ -34,6 +37,7 @@
 from CTFd.utils.exports import import_ctf as import_ctf_util
 from CTFd.utils.helpers import get_errors
 from CTFd.utils.security.auth import logout_user
+from CTFd.utils.uploads import delete_file
 from CTFd.utils.user import is_admin
 
 admin = Blueprint("admin", __name__)
@@ -176,19 +180,60 @@ def config():
 @admins_only
 def reset():
     if request.method == "POST":
-        # Truncate Users, Teams, Submissions, Solves, Notifications, Awards, Unlocks, Tracking
-        Tracking.query.delete()
-        Solves.query.delete()
-        Submissions.query.delete()
-        Awards.query.delete()
-        Unlocks.query.delete()
-        Users.query.delete()
-        Teams.query.delete()
-        set_config("setup", False)
+        require_setup = False
+        logout = False
+        next_url = url_for("admin.statistics")
+
+        data = request.form
+
+        if data.get("pages"):
+            _pages = Pages.query.all()
+            for p in _pages:
+                for f in p.files:
+                    delete_file(file_id=f.id)
+
+            Pages.query.delete()
+
+        if data.get("notifications"):
+            Notifications.query.delete()
+
+        if data.get("challenges"):
+            _challenges = Challenges.query.all()
+            for c in _challenges:
+                for f in c.files:
+                    delete_file(file_id=f.id)
+            Challenges.query.delete()
+
+        if data.get("accounts"):
+            Users.query.delete()
+            Teams.query.delete()
+            require_setup = True
+            logout = True
+
+        if data.get("submissions"):
+            Solves.query.delete()
+            Submissions.query.delete()
+            Awards.query.delete()
+            Unlocks.query.delete()
+            Tracking.query.delete()
+
+        if require_setup:
+            set_config("setup", False)
+            cache.clear()
+            logout_user()
+            next_url = url_for("views.setup")
+
         db.session.commit()
-        cache.clear()
-        logout_user()
+
+        clear_pages()
+        clear_standings()
+        clear_config()
+
+        if logout is True:
+            cache.clear()
+            logout_user()
+
         db.session.close()
-        return redirect(url_for("views.setup"))
+        return redirect(next_url)
 
     return render_template("admin/reset.html")
diff --git a/populate.py b/populate.py
--- a/populate.py
+++ b/populate.py
@@ -7,6 +7,7 @@
 import argparse
 
 from CTFd import create_app
+from CTFd.cache import clear_config, clear_standings, clear_pages
 from CTFd.models import (
     Users,
     Teams,
@@ -338,3 +339,7 @@ def random_chance():
 
         db.session.commit()
         db.session.close()
+
+        clear_config()
+        clear_standings()
+        clear_pages()


Test patch:
diff --git a/tests/admin/test_config.py b/tests/admin/test_config.py
--- a/tests/admin/test_config.py
+++ b/tests/admin/test_config.py
@@ -1,19 +1,36 @@
 import random
 
-from CTFd.models import Challenges, Fails, Solves, Teams, Tracking, Users
+from CTFd.models import (
+    Awards,
+    Challenges,
+    Fails,
+    Files,
+    Flags,
+    Hints,
+    Notifications,
+    Pages,
+    Solves,
+    Submissions,
+    Tags,
+    Teams,
+    Tracking,
+    Unlocks,
+    Users,
+)
 from tests.helpers import (
     create_ctfd,
     destroy_ctfd,
     gen_award,
     gen_challenge,
     gen_fail,
+    gen_file,
     gen_flag,
+    gen_hint,
     gen_solve,
     gen_team,
     gen_tracking,
     gen_user,
     login_as_user,
-    register_user,
 )
 
 
@@ -25,6 +42,12 @@ def test_reset():
         for x in range(10):
             chal = gen_challenge(app.db, name="chal_name{}".format(x))
             gen_flag(app.db, challenge_id=chal.id, content="flag")
+            gen_hint(app.db, challenge_id=chal.id)
+            gen_file(
+                app.db,
+                location="{name}/{name}.file".format(name=chal.name),
+                challenge_id=chal.id,
+            )
 
         for x in range(10):
             user = base_user + str(x)
@@ -35,18 +58,77 @@ def test_reset():
             gen_fail(app.db, user_id=user_obj.id, challenge_id=random.randint(1, 10))
             gen_tracking(app.db, user_id=user_obj.id)
 
+        # Add PageFiles
+        for x in range(5):
+            gen_file(
+                app.db,
+                location="page_file{name}/page_file{name}.file".format(name=x),
+                page_id=1,
+            )
+
         assert Users.query.count() == 11  # 11 because of the first admin user
         assert Challenges.query.count() == 10
+        assert (
+            Files.query.count() == 15
+        )  # This should be 11 because ChallengeFiles=10 and PageFiles=5
+        assert Flags.query.count() == 10
+        assert Hints.query.count() == 10
+        assert Submissions.query.count() == 20
+        assert Pages.query.count() == 1
+        assert Tracking.query.count() == 10
 
-        register_user(app)
         client = login_as_user(app, name="admin", password="password")
 
         with client.session_transaction() as sess:
-            data = {"nonce": sess.get("nonce")}
-            client.post("/admin/reset", data=data)
+            data = {"nonce": sess.get("nonce"), "pages": "on"}
+            r = client.post("/admin/reset", data=data)
+            assert r.location.endswith("/admin/statistics")
+        assert Pages.query.count() == 0
+        assert Users.query.count() == 11
+        assert Challenges.query.count() == 10
+        assert Tracking.query.count() == 11
+        assert Files.query.count() == 10
 
-        assert Users.query.count() == 0
+        with client.session_transaction() as sess:
+            data = {"nonce": sess.get("nonce"), "notifications": "on"}
+            r = client.post("/admin/reset", data=data)
+            assert r.location.endswith("/admin/statistics")
+        assert Notifications.query.count() == 0
+        assert Users.query.count() == 11
         assert Challenges.query.count() == 10
+        assert Tracking.query.count() == 11
+
+        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 Flags.query.count() == 0
+        assert Hints.query.count() == 0
+        assert Files.query.count() == 0
+        assert Tags.query.count() == 0
+        assert Users.query.count() == 11
+        assert Tracking.query.count() == 11
+
+        with client.session_transaction() as sess:
+            data = {"nonce": sess.get("nonce"), "submissions": "on"}
+            r = client.post("/admin/reset", data=data)
+            assert r.location.endswith("/admin/statistics")
+        assert Submissions.query.count() == 0
+        assert Solves.query.count() == 0
+        assert Fails.query.count() == 0
+        assert Awards.query.count() == 0
+        assert Unlocks.query.count() == 0
+        assert Users.query.count() == 11
+        assert Challenges.query.count() == 0
+        assert Flags.query.count() == 0
+        assert Tracking.query.count() == 0
+
+        with client.session_transaction() as sess:
+            data = {"nonce": sess.get("nonce"), "accounts": "on"}
+            r = client.post("/admin/reset", data=data)
+            assert r.location.endswith("/setup")
+        assert Users.query.count() == 0
         assert Solves.query.count() == 0
         assert Fails.query.count() == 0
         assert Tracking.query.count() == 0
@@ -62,6 +144,12 @@ def test_reset_team_mode():
         for x in range(10):
             chal = gen_challenge(app.db, name="chal_name{}".format(x))
             gen_flag(app.db, challenge_id=chal.id, content="flag")
+            gen_hint(app.db, challenge_id=chal.id)
+            gen_file(
+                app.db,
+                location="{name}/{name}.file".format(name=chal.name),
+                challenge_id=chal.id,
+            )
 
         for x in range(10):
             user = base_user + str(x)
@@ -78,22 +166,85 @@ def test_reset_team_mode():
             gen_fail(app.db, user_id=user_obj.id, challenge_id=random.randint(1, 10))
             gen_tracking(app.db, user_id=user_obj.id)
 
+        # Add PageFiles
+        for x in range(5):
+            gen_file(
+                app.db,
+                location="page_file{name}/page_file{name}.file".format(name=x),
+                page_id=1,
+            )
+
         assert Teams.query.count() == 10
-        assert (
-            Users.query.count() == 51
-        )  # 10 random users, 40 users (10 teams * 4), 1 admin user
+        # 10 random users, 40 users (10 teams * 4), 1 admin user
+        assert Users.query.count() == 51
         assert Challenges.query.count() == 10
+        assert (
+            Files.query.count() == 15
+        )  # This should be 11 because ChallengeFiles=10 and PageFiles=5
+        assert Flags.query.count() == 10
+        assert Hints.query.count() == 10
+        assert Submissions.query.count() == 20
+        assert Solves.query.count() == 10
+        assert Fails.query.count() == 10
+        assert Tracking.query.count() == 10
 
-        register_user(app)
         client = login_as_user(app, name="admin", password="password")
 
         with client.session_transaction() as sess:
-            data = {"nonce": sess.get("nonce")}
-            client.post("/admin/reset", data=data)
+            data = {"nonce": sess.get("nonce"), "pages": "on"}
+            r = client.post("/admin/reset", data=data)
+            assert r.location.endswith("/admin/statistics")
+        assert Pages.query.count() == 0
+        assert Teams.query.count() == 10
+        assert Users.query.count() == 51
+        assert Challenges.query.count() == 10
+        assert Tracking.query.count() == 11
+        assert Files.query.count() == 10
 
-        assert Teams.query.count() == 0
-        assert Users.query.count() == 0
+        with client.session_transaction() as sess:
+            data = {"nonce": sess.get("nonce"), "notifications": "on"}
+            r = client.post("/admin/reset", data=data)
+            assert r.location.endswith("/admin/statistics")
+        assert Notifications.query.count() == 0
+        assert Teams.query.count() == 10
+        assert Users.query.count() == 51
         assert Challenges.query.count() == 10
+        assert Tracking.query.count() == 11
+
+        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 Flags.query.count() == 0
+        assert Hints.query.count() == 0
+        assert Files.query.count() == 0
+        assert Tags.query.count() == 0
+        assert Teams.query.count() == 10
+        assert Users.query.count() == 51
+        assert Tracking.query.count() == 11
+
+        with client.session_transaction() as sess:
+            data = {"nonce": sess.get("nonce"), "submissions": "on"}
+            r = client.post("/admin/reset", data=data)
+            assert r.location.endswith("/admin/statistics")
+        assert Submissions.query.count() == 0
+        assert Solves.query.count() == 0
+        assert Fails.query.count() == 0
+        assert Awards.query.count() == 0
+        assert Unlocks.query.count() == 0
+        assert Teams.query.count() == 10
+        assert Users.query.count() == 51
+        assert Challenges.query.count() == 0
+        assert Flags.query.count() == 0
+        assert Tracking.query.count() == 0
+
+        with client.session_transaction() as sess:
+            data = {"nonce": sess.get("nonce"), "accounts": "on"}
+            r = client.post("/admin/reset", data=data)
+            assert r.location.endswith("/setup")
+        assert Users.query.count() == 0
+        assert Teams.query.count() == 0
         assert Solves.query.count() == 0
         assert Fails.query.count() == 0
         assert Tracking.query.count() == 0