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

@@ -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"))

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