From 7e6d56694e12cde7b827c848d92c01a35d35314c Mon Sep 17 00:00:00 2001
From: Kevin Chung <kchung@nyu.edu>
Date: Mon, 14 Aug 2017 20:27:54 -0400
Subject: [PATCH] Give challenge plugins the ability to specify the response
 message (#351)

---
 CTFd/challenges.py                  | 39 ++++++++++++++++-------------
 CTFd/plugins/challenges/__init__.py |  4 +--
 2 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/CTFd/challenges.py b/CTFd/challenges.py
index 31cb6590..8a120b5d 100644
--- a/CTFd/challenges.py
+++ b/CTFd/challenges.py
@@ -287,30 +287,33 @@ def chal(chalid):
                 })
 
             chal_class = get_chal_class(chal.type)
-            if chal_class.solve(chal, provided_key):
+            status, message = chal_class.solve(chal, provided_key)
+            if status:  # The challenge plugin says the input is right
                 if utils.ctftime():
                     solve = Solves(teamid=session['id'], chalid=chalid, ip=utils.get_ip(), flag=provided_key)
                     db.session.add(solve)
                     db.session.commit()
                     db.session.close()
                 logger.info("[{0}] {1} submitted {2} with kpm {3} [CORRECT]".format(*data))
-                return jsonify({'status': 1, 'message': 'Correct'})
-
-            if utils.ctftime():
-                wrong = WrongKeys(teamid=session['id'], chalid=chalid, ip=utils.get_ip(), flag=provided_key)
-                db.session.add(wrong)
-                db.session.commit()
-                db.session.close()
-            logger.info("[{0}] {1} submitted {2} with kpm {3} [WRONG]".format(*data))
-            # return '0' # key was wrong
-            if max_tries:
-                attempts_left = max_tries - fails - 1  # Off by one since fails has changed since it was gotten
-                tries_str = 'tries'
-                if attempts_left == 1:
-                    tries_str = 'try'
-                return jsonify({'status': 0, 'message': 'Incorrect. You have {} {} remaining.'.format(attempts_left, tries_str)})
-            else:
-                return jsonify({'status': 0, 'message': 'Incorrect'})
+                return jsonify({'status': 1, 'message': message})
+            else:  # The challenge plugin says the input is wrong
+                if utils.ctftime():
+                    wrong = WrongKeys(teamid=session['id'], chalid=chalid, ip=utils.get_ip(), flag=provided_key)
+                    db.session.add(wrong)
+                    db.session.commit()
+                    db.session.close()
+                logger.info("[{0}] {1} submitted {2} with kpm {3} [WRONG]".format(*data))
+                # return '0' # key was wrong
+                if max_tries:
+                    attempts_left = max_tries - fails - 1  # Off by one since fails has changed since it was gotten
+                    tries_str = 'tries'
+                    if attempts_left == 1:
+                        tries_str = 'try'
+                    if message[-1] not in '!().;?[]\{\}':  # Add a punctuation mark if there isn't one
+                        message = message + '.'
+                    return jsonify({'status': 0, 'message': '{} You have {} {} remaining.'.format(message, attempts_left, tries_str)})
+                else:
+                    return jsonify({'status': 0, 'message': message})
 
         # Challenge already solved
         else:
diff --git a/CTFd/plugins/challenges/__init__.py b/CTFd/plugins/challenges/__init__.py
index 2f00a8d6..60c93f1f 100644
--- a/CTFd/plugins/challenges/__init__.py
+++ b/CTFd/plugins/challenges/__init__.py
@@ -16,8 +16,8 @@ class CTFdStandardChallenge(BaseChallenge):
         chal_keys = Keys.query.filter_by(chal=chal.id).all()
         for chal_key in chal_keys:
             if get_key_class(chal_key.key_type).compare(chal_key.flag, provided_key):
-                return True
-        return False
+                return True, 'Correct'
+        return False, 'Incorrect'
 
 
 CHALLENGE_CLASSES = {