From 615ca5afe361343c61c31aaa99dab007ef9fce23 Mon Sep 17 00:00:00 2001
From: Nanne Baars <nanne.baars@owasp.org>
Date: Tue, 2 May 2017 03:25:31 +0200
Subject: [PATCH] Posting a flag shows a response in the UI (correct or
 incorrect)

---
 .../webgoat/assignments/AttackResult.java     |  6 +++++
 .../java/org/owasp/webgoat/plugin/Flag.java   | 23 +++++++++++--------
 .../resources/i18n/WebGoatLabels.properties   |  5 +++-
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/webgoat-container/src/main/java/org/owasp/webgoat/assignments/AttackResult.java b/webgoat-container/src/main/java/org/owasp/webgoat/assignments/AttackResult.java
index 4cf1dbad8..b10917d49 100644
--- a/webgoat-container/src/main/java/org/owasp/webgoat/assignments/AttackResult.java
+++ b/webgoat-container/src/main/java/org/owasp/webgoat/assignments/AttackResult.java
@@ -51,6 +51,12 @@ public class AttackResult {
             return this;
         }
 
+        public AttackResultBuilder lessonCompleted(boolean lessonCompleted, String resourceBundleKey) {
+            this.lessonCompleted = lessonCompleted;
+            this.feedbackResourceBundleKey = resourceBundleKey;
+            return this;
+        }
+
         public AttackResultBuilder feedbackArgs(Object... args) {
             this.feedbackArgs = args;
             return this;
diff --git a/webgoat-lessons/challenge/src/main/java/org/owasp/webgoat/plugin/Flag.java b/webgoat-lessons/challenge/src/main/java/org/owasp/webgoat/plugin/Flag.java
index 1cf13301c..944dba4ef 100644
--- a/webgoat-lessons/challenge/src/main/java/org/owasp/webgoat/plugin/Flag.java
+++ b/webgoat-lessons/challenge/src/main/java/org/owasp/webgoat/plugin/Flag.java
@@ -4,17 +4,18 @@ import com.google.common.collect.Maps;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
+import org.owasp.webgoat.assignments.AttackResult;
 import org.owasp.webgoat.assignments.Endpoint;
+import org.owasp.webgoat.i18n.PluginMessages;
+import org.owasp.webgoat.session.WebSession;
 import org.owasp.webgoat.users.UserTracker;
 import org.owasp.webgoat.users.UserTrackerRepository;
-import org.owasp.webgoat.session.WebSession;
 import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestMethod;
 import org.springframework.web.bind.annotation.RequestParam;
-import org.springframework.web.bind.annotation.ResponseStatus;
+import org.springframework.web.bind.annotation.ResponseBody;
 
 import javax.annotation.PostConstruct;
 import java.util.Map;
@@ -33,6 +34,8 @@ public class Flag extends Endpoint {
     private UserTrackerRepository userTrackerRepository;
     @Autowired
     private WebSession webSession;
+    @Autowired
+    private PluginMessages pluginMessages;
 
     @AllArgsConstructor
     private class FlagPosted {
@@ -43,7 +46,7 @@ public class Flag extends Endpoint {
     @PostConstruct
     public void initFlags() {
         IntStream.range(1, 6).forEach(i -> FLAGS.put(i, UUID.randomUUID().toString()));
-        FLAGS.entrySet().stream().forEach(e ->  log.debug("Flag {} {}", e.getKey(), e.getValue()));
+        FLAGS.entrySet().stream().forEach(e -> log.debug("Flag {} {}", e.getKey(), e.getValue()));
     }
 
     @Override
@@ -52,21 +55,21 @@ public class Flag extends Endpoint {
     }
 
     @RequestMapping(method = RequestMethod.POST, produces = MediaType.APPLICATION_JSON_VALUE)
-    @ResponseStatus(HttpStatus.OK)
-    public void postFlag(@RequestParam String flag) {
+    @ResponseBody
+    public AttackResult postFlag(@RequestParam String flag) {
         UserTracker userTracker = userTrackerRepository.findOne(webSession.getUserName());
-        if (userTracker == null) {
-            userTracker = new UserTracker(webSession.getUserName());
-        }
         String currentChallenge = webSession.getCurrentLesson().getName();
         int challengeNumber = Integer.valueOf(currentChallenge.substring(currentChallenge.length() - 1, currentChallenge.length()));
         String expectedFlag = FLAGS.get(challengeNumber);
+        final AttackResult attackResult;
         if (expectedFlag.equals(flag)) {
             userTracker.assignmentSolved(webSession.getCurrentLesson(), "Assignment" + challengeNumber);
+            attackResult = new AttackResult.AttackResultBuilder(pluginMessages).lessonCompleted(true, "challenge.flag.correct").build();
         } else {
             userTracker.assignmentFailed(webSession.getCurrentLesson());
+            attackResult = new AttackResult.AttackResultBuilder(pluginMessages).feedback("challenge.flag.incorrect").build();
         }
         userTrackerRepository.save(userTracker);
+        return attackResult;
     }
-
 }
diff --git a/webgoat-lessons/challenge/src/main/resources/i18n/WebGoatLabels.properties b/webgoat-lessons/challenge/src/main/resources/i18n/WebGoatLabels.properties
index 3a425339e..9e5284d90 100644
--- a/webgoat-lessons/challenge/src/main/resources/i18n/WebGoatLabels.properties
+++ b/webgoat-lessons/challenge/src/main/resources/i18n/WebGoatLabels.properties
@@ -9,4 +9,7 @@ challenge.close=This is not the correct password for tom, please try again.
 
 user.exists=User {0} already exists please try to register with a different username.
 user.created=User {0} created, please proceed to the login page.
-input.invalid=Input for user, email and/or password is empty or too long, please fill in all field and/or limit all fields to 30 characters.
\ No newline at end of file
+input.invalid=Input for user, email and/or password is empty or too long, please fill in all field and/or limit all fields to 30 characters.
+
+challenge.flag.correct=Congratulations you have solved the challenge!!
+challenge.flag.incorrect=Sorry this is not the correct flag, please try again.
\ No newline at end of file