From 1b372b46be5396f39433031f257ef8ffe19503a9 Mon Sep 17 00:00:00 2001 From: l4cr0ss Date: Sat, 6 Jun 2020 10:08:26 -0400 Subject: [PATCH] Incorporate PR feedback from kileak + Revert changes made in b88b8cc to the util function `get_challenge_from_args` + Add method `get_challenge_by_args_or_channel` for tag commands The bug in the `!solve` command was introduced as a result of the utility function changes that are undone here, fixing the bug. --- handlers/challenge_handler.py | 40 +++++++++++++++++------------------ util/util.py | 32 +++++++++++++++++++++++++--- 2 files changed, 49 insertions(+), 23 deletions(-) diff --git a/handlers/challenge_handler.py b/handlers/challenge_handler.py index f389fac..06432a2 100644 --- a/handlers/challenge_handler.py +++ b/handlers/challenge_handler.py @@ -22,7 +22,7 @@ class AddChallengeTagCommand(Command): def execute(cls, slack_wrapper, args, timestamp, channel_id, user_id, user_is_admin): tags = None - challenge = get_challenge_from_args(ChallengeHandler.DB, args, channel_id) + challenge = get_challenge_from_args_or_channel(ChallengeHandler.DB, args, channel_id) if challenge.channel_id == channel_id: # We were called from the Challenge channel @@ -50,7 +50,7 @@ class RemoveChallengeTagCommand(Command): def execute(cls, slack_wrapper, args, timestamp, channel_id, user_id, user_is_admin): tags = None - challenge = get_challenge_from_args(ChallengeHandler.DB, args, channel_id) + challenge = get_challenge_from_args_or_channel(ChallengeHandler.DB, args, channel_id) if challenge.channel_id == channel_id: # We were called from the Challenge channel @@ -603,24 +603,24 @@ class SolveCommand(Command): @classmethod def execute(cls, slack_wrapper, args, timestamp, channel_id, user_id, user_is_admin): """Execute the Solve command.""" + challenge = "" - # Get the challenge from the channel id or the first argument to !solve. - challenge = get_challenge_from_args(ChallengeHandler.DB, args, channel_id) - if not challenge: - raise InvalidCommand("You must be in a CTF or Challenge channel to use this command.") + if args: + challenge = get_challenge_from_args(ChallengeHandler.DB, args, channel_id) - if channel_id == challenge.channel_id: - # !solve called from the challenge - additional args are - # assumed to be other solvers. - additional_args = args if args else [] - elif channel_id == challenge.ctf_channel_id: - # !solve called from the ctf - first argument must be the name - # of the challenge. - additional_args = args[1:] if len(args) > 1 else [] + if not challenge: + challenge = get_challenge_by_channel_id(ChallengeHandler.DB, channel_id) + additional_args = args if args else [] + else: + additional_args = args[1:] if len(args) > 1 else [] else: - # The user made a typo, or is perhaps confused. - # NOTE: This branch should be unreachable. - raise InvalidCommand("You must be in a CTF or Challenge channel to use this command.") + # No arguments => direct way of resolving challenge + challenge = get_challenge_by_channel_id(ChallengeHandler.DB, channel_id) + + additional_args = [] + + if not challenge: + raise InvalidCommand("This challenge does not exist.") additional_solver = [] @@ -635,9 +635,9 @@ def execute(cls, slack_wrapper, args, timestamp, channel_id, user_id, user_is_ad if user_obj['ok']: add_solve = get_display_name(user_obj) - if add_solve not in solver_list: - solver_list.append(add_solve) - additional_solver.append(add_solve) + if add_solve not in solver_list: + solver_list.append(add_solve) + additional_solver.append(add_solve) # Update database ctfs = pickle.load(open(ChallengeHandler.DB, "rb")) diff --git a/util/util.py b/util/util.py index 096f976..fe4dfc7 100644 --- a/util/util.py +++ b/util/util.py @@ -119,9 +119,13 @@ def get_challenge_by_name(database, challenge_name, ctf_channel_id): return None -def get_challenge_from_args(database, args, channel_id): +def get_challenge_from_args_or_channel(database, args, channel_id): """ Helper method for getting a Challenge either from arguments or current channel. + Return the corresponding Challenge if called from a challenge channel. + Return the Challenge correpsonding to the first argument if called from the + CTF channel. + Return None if no Challenge can be found. """ # Check if we're currently in a challenge channel @@ -131,8 +135,7 @@ def get_challenge_from_args(database, args, channel_id): # User is in the challenge channel challenge = current_chal else: - # User is in the ctf channel => Check for challenge by name in - # current challenge + # Assume user is in the ctf channel try: challenge_name = args[0].lower().strip("*") challenge = get_challenge_by_name(database, challenge_name, channel_id) @@ -142,6 +145,29 @@ def get_challenge_from_args(database, args, channel_id): return challenge +def get_challenge_from_args(database, args, channel_id): + """ + Helper method for getting a Challenge either from arguments or current channel. + Return None if called from a Challenge channel. + """ + # Multiple arguments: Need to check if a challenge was specified or not + challenge_name = args[0].lower().strip("*") + + # Check if we're currently in a challenge channel + current_chal = get_challenge_by_channel_id(database, channel_id) + + if current_chal: + # User is in a challenge channel => Check for challenge by name + # in parent ctf channel + challenge = get_challenge_by_name(database, challenge_name, current_chal.ctf_channel_id) + else: + # User is in the ctf channel => Check for challenge by name in + # current challenge + challenge = get_challenge_by_name(database, challenge_name, channel_id) + + return challenge + + def get_challenge_by_channel_id(database, challenge_channel_id): """ Fetch a Challenge object in the database with a given channel ID