refactor: update challenge code

- Flags are now wired through a Spring config - Introduced Flag class - Removed Flags from the FlagController
nklkarthi · Feb 22, 2023 · 5dbe2ea · 5dbe2ea
1 parent 1b49b2f
commit 5dbe2ea
Show file tree

Hide file tree

Showing 9 changed files with 131 additions and 145 deletions.
diff --git a/src/main/java/org/owasp/webgoat/lessons/challenges/Flag.java b/src/main/java/org/owasp/webgoat/lessons/challenges/Flag.java
@@ -1,89 +1,13 @@
-/*
- * This file is part of WebGoat, an Open Web Application Security Project utility. For details, please see http://www.owasp.org/
- *
- * Copyright (c) 2002 - 2019 Bruce Mayhew
- *
- * This program is free software; you can redistribute it and/or modify it under the terms of the
- * GNU General Public License as published by the Free Software Foundation; either version 2 of the
- * License, or (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without
- * even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
- * General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License along with this program; if
- * not, write to the Free Software Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA
- * 02111-1307, USA.
- *
- * Getting Source ==============
- *
- * Source for this application is maintained at https://github.com/WebGoat/WebGoat, a repository for free software projects.
- */
-
 package org.owasp.webgoat.lessons.challenges;
 
-import java.util.HashMap;
-import java.util.Map;
-import java.util.UUID;
-import java.util.stream.IntStream;
-import javax.annotation.PostConstruct;
-import lombok.AllArgsConstructor;
-import lombok.Getter;
-import org.owasp.webgoat.container.assignments.AssignmentEndpoint;
-import org.owasp.webgoat.container.assignments.AttackResult;
-import org.owasp.webgoat.container.session.WebSession;
-import org.owasp.webgoat.container.users.UserTracker;
-import org.owasp.webgoat.container.users.UserTrackerRepository;
-import org.springframework.beans.factory.annotation.Autowired;
-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.ResponseBody;
-import org.springframework.web.bind.annotation.RestController;
-
-/**
- * @author nbaars
- * @since 3/23/17.
- */
-@RestController
-public class Flag extends AssignmentEndpoint {
-
-  public static final Map<Integer, String> FLAGS = new HashMap<>();
-  @Autowired private UserTrackerRepository userTrackerRepository;
-  @Autowired private WebSession webSession;
-
-  @AllArgsConstructor
-  private class FlagPosted {
-    @Getter private boolean lessonCompleted;
-  }
+public record Flag(int number, String answer) {
 
-  @PostConstruct
-  public void initFlags() {
-    IntStream.range(1, 10).forEach(i -> FLAGS.put(i, UUID.randomUUID().toString()));
+  public boolean isCorrect(String flag) {
+    return answer.equals(flag);
   }
 
-  @RequestMapping(
-      path = "/challenge/flag",
-      method = RequestMethod.POST,
-      produces = MediaType.APPLICATION_JSON_VALUE)
-  @ResponseBody
-  public AttackResult postFlag(@RequestParam String flag) {
-    UserTracker userTracker = userTrackerRepository.findByUser(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 = success(this).feedback("challenge.flag.correct").build();
-    } else {
-      userTracker.assignmentFailed(webSession.getCurrentLesson());
-      attackResult = failed(this).feedback("challenge.flag.incorrect").build();
-    }
-    userTrackerRepository.save(userTracker);
-    return attackResult;
+  @Override
+  public String toString() {
+    return answer;
   }
 }
diff --git a/src/main/java/org/owasp/webgoat/lessons/challenges/FlagController.java b/src/main/java/org/owasp/webgoat/lessons/challenges/FlagController.java
@@ -0,0 +1,62 @@
+/*
+ * This file is part of WebGoat, an Open Web Application Security Project utility. For details, please see http://www.owasp.org/
+ *
+ * Copyright (c) 2002 - 2019 Bruce Mayhew
+ *
+ * This program is free software; you can redistribute it and/or modify it under the terms of the
+ * GNU General Public License as published by the Free Software Foundation; either version 2 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without
+ * even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License along with this program; if
+ * not, write to the Free Software Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA
+ * 02111-1307, USA.
+ *
+ * Getting Source ==============
+ *
+ * Source for this application is maintained at https://github.com/WebGoat/WebGoat, a repository for free software projects.
+ */
+
+package org.owasp.webgoat.lessons.challenges;
+
+import lombok.AllArgsConstructor;
+import org.owasp.webgoat.container.assignments.AssignmentEndpoint;
+import org.owasp.webgoat.container.assignments.AttackResult;
+import org.owasp.webgoat.container.session.WebSession;
+import org.owasp.webgoat.container.users.UserTracker;
+import org.owasp.webgoat.container.users.UserTrackerRepository;
+import org.springframework.http.MediaType;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestParam;
+import org.springframework.web.bind.annotation.ResponseBody;
+import org.springframework.web.bind.annotation.RestController;
+
+@RestController
+@AllArgsConstructor
+public class FlagController extends AssignmentEndpoint {
+
+  private final UserTrackerRepository userTrackerRepository;
+  private final WebSession webSession;
+  private final Flags flags;
+
+  @PostMapping(path = "/challenge/flag", produces = MediaType.APPLICATION_JSON_VALUE)
+  @ResponseBody
+  public AttackResult postFlag(@RequestParam String flag) {
+    UserTracker userTracker = userTrackerRepository.findByUser(webSession.getUserName());
+    Flag expectedFlag = flags.getFlag(webSession.getCurrentLesson());
+    final AttackResult attackResult;
+    if (expectedFlag.isCorrect(flag)) {
+      userTracker.assignmentSolved(
+          webSession.getCurrentLesson(), "Assignment" + expectedFlag.number());
+      attackResult = success(this).feedback("challenge.flag.correct").build();
+    } else {
+      userTracker.assignmentFailed(webSession.getCurrentLesson());
+      attackResult = failed(this).feedback("challenge.flag.incorrect").build();
+    }
+    userTrackerRepository.save(userTracker);
+    return attackResult;
+  }
+}
diff --git a/src/main/java/org/owasp/webgoat/lessons/challenges/Flags.java b/src/main/java/org/owasp/webgoat/lessons/challenges/Flags.java
@@ -0,0 +1,27 @@
+package org.owasp.webgoat.lessons.challenges;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+import java.util.stream.IntStream;
+import org.owasp.webgoat.container.lessons.Lesson;
+import org.springframework.context.annotation.Configuration;
+
+@Configuration
+public class Flags {
+  private final Map<Integer, Flag> FLAGS = new HashMap<>();
+
+  public Flags() {
+    IntStream.range(1, 10).forEach(i -> FLAGS.put(i, new Flag(i, UUID.randomUUID().toString())));
+  }
+
+  public Flag getFlag(Lesson forLesson) {
+    String lessonName = forLesson.getName();
+    int challengeNumber = Integer.valueOf(lessonName.substring(lessonName.length() - 1));
+    return FLAGS.get(challengeNumber);
+  }
+
+  public Flag getFlag(int flagNumber) {
+    return FLAGS.get(flagNumber);
+  }
+}
diff --git a/src/main/java/org/owasp/webgoat/lessons/challenges/challenge1/Assignment1.java b/src/main/java/org/owasp/webgoat/lessons/challenges/challenge1/Assignment1.java
@@ -2,11 +2,10 @@
 
 import static org.owasp.webgoat.lessons.challenges.SolutionConstants.PASSWORD;
 
-import javax.servlet.http.HttpServletRequest;
+import lombok.RequiredArgsConstructor;
 import org.owasp.webgoat.container.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.container.assignments.AttackResult;
-import org.owasp.webgoat.lessons.challenges.Flag;
-import org.springframework.util.StringUtils;
+import org.owasp.webgoat.lessons.challenges.Flags;
 import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
@@ -43,27 +42,25 @@
  * @since August 11, 2016
  */
 @RestController
+@RequiredArgsConstructor
 public class Assignment1 extends AssignmentEndpoint {
 
+  private final Flags flags;
+
   @PostMapping("/challenge/1")
   @ResponseBody
-  public AttackResult completed(
-      @RequestParam String username, @RequestParam String password, HttpServletRequest request) {
+  public AttackResult completed(@RequestParam String username, @RequestParam String password) {
     boolean ipAddressKnown = true;
     boolean passwordCorrect =
         "admin".equals(username)
             && PASSWORD
                 .replace("1234", String.format("%04d", ImageServlet.PINCODE))
                 .equals(password);
     if (passwordCorrect && ipAddressKnown) {
-      return success(this).feedback("challenge.solved").feedbackArgs(Flag.FLAGS.get(1)).build();
+      return success(this).feedback("challenge.solved").feedbackArgs(flags.getFlag(1)).build();
     } else if (passwordCorrect) {
       return failed(this).feedback("ip.address.unknown").build();
     }
     return failed(this).build();
   }
-
-  public static boolean containsHeader(HttpServletRequest request) {
-    return StringUtils.hasText(request.getHeader("X-Forwarded-For"));
-  }
 }
diff --git a/src/main/java/org/owasp/webgoat/lessons/challenges/challenge1/ImageServlet.java b/src/main/java/org/owasp/webgoat/lessons/challenges/challenge1/ImageServlet.java
@@ -4,19 +4,17 @@
 import static org.springframework.web.bind.annotation.RequestMethod.POST;
 
 import java.io.IOException;
-import java.security.SecureRandom;
-import javax.servlet.http.HttpServlet;
+import java.util.Random;
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.http.MediaType;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
 @RestController
-public class ImageServlet extends HttpServlet {
+public class ImageServlet {
 
-  private static final long serialVersionUID = 9132775506936676850L;
-  public static final int PINCODE = new SecureRandom().nextInt(10000);
+  public static final int PINCODE = new Random().nextInt(10000);
 
   @RequestMapping(
       method = {GET, POST},

diff --git a/src/main/java/org/owasp/webgoat/lessons/challenges/challenge5/Assignment5.java b/src/main/java/org/owasp/webgoat/lessons/challenges/challenge5/Assignment5.java
@@ -24,11 +24,12 @@
 
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
+import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.owasp.webgoat.container.LessonDataSource;
 import org.owasp.webgoat.container.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.container.assignments.AttackResult;
-import org.owasp.webgoat.lessons.challenges.Flag;
+import org.owasp.webgoat.lessons.challenges.Flags;
 import org.springframework.util.StringUtils;
 import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestParam;
@@ -37,13 +38,11 @@
 
 @RestController
 @Slf4j
+@RequiredArgsConstructor
 public class Assignment5 extends AssignmentEndpoint {
 
   private final LessonDataSource dataSource;
-
-  public Assignment5(LessonDataSource dataSource) {
-    this.dataSource = dataSource;
-  }
+  private final Flags flags;
 
   @PostMapping("/challenge/5")
   @ResponseBody
@@ -66,7 +65,7 @@ public AttackResult login(
       ResultSet resultSet = statement.executeQuery();
 
       if (resultSet.next()) {
-        return success(this).feedback("challenge.solved").feedbackArgs(Flag.FLAGS.get(5)).build();
+        return success(this).feedback("challenge.solved").feedbackArgs(flags.getFlag(5)).build();
       } else {
         return failed(this).feedback("challenge.close").build();
       }

diff --git a/src/main/java/org/owasp/webgoat/lessons/challenges/challenge7/Assignment7.java b/src/main/java/org/owasp/webgoat/lessons/challenges/challenge7/Assignment7.java
@@ -8,9 +8,8 @@
 import org.owasp.webgoat.container.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.container.assignments.AttackResult;
 import org.owasp.webgoat.lessons.challenges.Email;
-import org.owasp.webgoat.lessons.challenges.Flag;
+import org.owasp.webgoat.lessons.challenges.Flags;
 import org.owasp.webgoat.lessons.challenges.SolutionConstants;
-import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.http.HttpStatus;
@@ -44,10 +43,16 @@ public class Assignment7 extends AssignmentEndpoint {
           + "Kind regards, \n"
           + "Team WebGoat";
 
-  @Autowired private RestTemplate restTemplate;
+  private final Flags flags;
+  private final RestTemplate restTemplate;
+  private final String webWolfMailURL;
 
-  @Value("${webwolf.mail.url}")
-  private String webWolfMailURL;
+  public Assignment7(
+      Flags flags, RestTemplate restTemplate, @Value("${webwolf.mail.url}") String webWolfMailURL) {
+    this.flags = flags;
+    this.restTemplate = restTemplate;
+    this.webWolfMailURL = webWolfMailURL;
+  }
 
   @GetMapping("/challenge/7/reset-password/{link}")
   public ResponseEntity<String> resetPassword(@PathVariable(value = "link") String link) {
@@ -58,7 +63,7 @@ public ResponseEntity<String> resetPassword(@PathVariable(value = "link") String
                   + "<img src='/WebGoat/images/hi-five-cat.jpg'>"
                   + "<br/><br/>Here is your flag: "
                   + "<b>"
-                  + Flag.FLAGS.get(7)
+                  + flags.getFlag(7)
                   + "</b>");
     }
     return ResponseEntity.status(HttpStatus.I_AM_A_TEAPOT)

diff --git a/src/main/java/org/owasp/webgoat/lessons/challenges/challenge8/Assignment8.java b/src/main/java/org/owasp/webgoat/lessons/challenges/challenge8/Assignment8.java
@@ -4,23 +4,21 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import javax.servlet.http.HttpServletRequest;
+import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.owasp.webgoat.container.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.container.assignments.AttackResult;
-import org.owasp.webgoat.lessons.challenges.Flag;
+import org.owasp.webgoat.lessons.challenges.Flags;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.PathVariable;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-/**
- * @author nbaars
- * @since 4/8/17.
- */
 @RestController
 @Slf4j
+@RequiredArgsConstructor
 public class Assignment8 extends AssignmentEndpoint {
 
   private static final Map<Integer, Integer> votes = new HashMap<>();
@@ -33,6 +31,8 @@ public class Assignment8 extends AssignmentEndpoint {
     votes.put(5, 300);
   }
 
+  private final Flags flags;
+
   @GetMapping(value = "/challenge/8/vote/{stars}", produces = MediaType.APPLICATION_JSON_VALUE)
   @ResponseBody
   public ResponseEntity<?> vote(
@@ -47,7 +47,7 @@ public ResponseEntity<?> vote(
     Integer allVotesForStar = votes.getOrDefault(nrOfStars, 0);
     votes.put(nrOfStars, allVotesForStar + 1);
     return ResponseEntity.ok()
-        .header("X-Flag", "Thanks for voting, your flag is: " + Flag.FLAGS.get(8))
+        .header("X-FlagController", "Thanks for voting, your flag is: " + flags.getFlag(8))
         .build();
   }