automatically grant belt awards going forward

dojo_plugin/__init__.py

@@ -15,7 +15,7 @@
 from CTFd.plugins.challenges import CHALLENGE_CLASSES, BaseChallenge
 from CTFd.plugins.flags import FLAG_CLASSES, BaseFlag, FlagException
 
-from .models import Dojos, DojoChallenges
+from .models import Dojos, DojoChallenges, Belts
 from .config import DOJO_HOST, bootstrap
 from .utils import unserialize_user_flag, render_markdown
 from .utils.discord import get_discord_user, get_discord_roles, add_role, send_message
@@ -43,18 +43,27 @@ class DojoChallenge(BaseChallenge):
     def solve(cls, user, team, challenge, request):
         super().solve(user, team, challenge, request)
 
+        current_belts = get_user_belts(user)
+        for belt in current_belts:
+            belt_award = Belts.query.filter_by(user=user, name=belt).one()
+            if belt_award:
+                continue
+            db.session.add(Belts(user=user, name=belt))
+            db.session.commit()
+
         discord_user = get_discord_user(user.id)
         if not discord_user:
             return
 
         discord_roles = get_discord_roles()
-        for belt in get_user_belts(user):
-            if discord_roles.get(belt) in discord_user["roles"]:
+        for belt in current_belts:
+            belt_role = belt.title() + " Belt"
+            if discord_roles.get(belt_role) in discord_user["roles"]:
                 continue
             user_mention = f"<@{discord_user['user']['id']}>"
-            message = f"{user_mention} earned their {belt}! :tada:"
+            message = f"{user_mention} earned their {belt_role}! :tada:"
             print(message, flush=True)
-            add_role(discord_user["user"]["id"], belt)
+            add_role(discord_user["user"]["id"], belt_role)
             send_message(message, "belting-ceremony")
 
 

dojo_plugin/utils/belts.py

@@ -23,7 +23,7 @@ def get_user_belts(user):
         dojo = Dojos.query.filter(Dojos.official, Dojos.id == dojo_id).one()
         if not dojo.completed(user):
             break
-        result.append(belt.title() + " Belt")
+        result.append(belt)
     return result
 
 @cache.memoize(timeout=60)