diff --git a/dojo_plugin/__init__.py b/dojo_plugin/__init__.py index 802d09cfd..fbb0fb561 100644 --- a/dojo_plugin/__init__.py +++ b/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") diff --git a/dojo_plugin/utils/belts.py b/dojo_plugin/utils/belts.py index e7a2bfb9d..abf5c28ba 100644 --- a/dojo_plugin/utils/belts.py +++ b/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)