From eb13ebc26fa35c319aeab9c003ca785ab6c67a77 Mon Sep 17 00:00:00 2001
From: Nanne Baars <nanne.baars@owasp.org>
Date: Sat, 15 Apr 2017 13:59:57 +0200
Subject: [PATCH] Assignments were not grouped per lesson in the same package

---
 .../src/main/java/org/owasp/webgoat/WebGoat.java           | 1 +
 .../java/org/owasp/webgoat/plugins/PluginResource.java     | 7 ++++---
 .../main/java/org/owasp/webgoat/plugins/PluginsLoader.java | 2 +-
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/webgoat-container/src/main/java/org/owasp/webgoat/WebGoat.java b/webgoat-container/src/main/java/org/owasp/webgoat/WebGoat.java
index 9e88dec97..36a2c761d 100644
--- a/webgoat-container/src/main/java/org/owasp/webgoat/WebGoat.java
+++ b/webgoat-container/src/main/java/org/owasp/webgoat/WebGoat.java
@@ -127,4 +127,5 @@ public class WebGoat extends SpringBootServletInitializer {
         }
     }
 
+
 }
diff --git a/webgoat-container/src/main/java/org/owasp/webgoat/plugins/PluginResource.java b/webgoat-container/src/main/java/org/owasp/webgoat/plugins/PluginResource.java
index 73e3ee03b..172b7b965 100644
--- a/webgoat-container/src/main/java/org/owasp/webgoat/plugins/PluginResource.java
+++ b/webgoat-container/src/main/java/org/owasp/webgoat/plugins/PluginResource.java
@@ -30,14 +30,15 @@ public class PluginResource {
     public List<Class<Endpoint>> getEndpoints() {
         return classes.stream().
                 filter(c -> c.getSuperclass() == AssignmentEndpoint.class || c.getSuperclass() == Endpoint.class).
-                map(c -> (Class<Endpoint>)c).
+                map(c -> (Class<Endpoint>) c).
                 collect(Collectors.toList());
     }
 
-    public List<Class<AssignmentEndpoint>> getAssignments() {
+    public List<Class<AssignmentEndpoint>> getAssignments(Class lesson) {
         return classes.stream().
                 filter(c -> c.getSuperclass() == AssignmentEndpoint.class).
-                map(c -> (Class<AssignmentEndpoint>)c).
+                filter(c -> c.getPackage().equals(lesson.getPackage())).
+                map(c -> (Class<AssignmentEndpoint>) c).
                 collect(Collectors.toList());
     }
 
diff --git a/webgoat-container/src/main/java/org/owasp/webgoat/plugins/PluginsLoader.java b/webgoat-container/src/main/java/org/owasp/webgoat/plugins/PluginsLoader.java
index 594b185ef..20e193025 100644
--- a/webgoat-container/src/main/java/org/owasp/webgoat/plugins/PluginsLoader.java
+++ b/webgoat-container/src/main/java/org/owasp/webgoat/plugins/PluginsLoader.java
@@ -74,7 +74,7 @@ public class PluginsLoader {
                     } catch (Exception e) {
                         log.error("Error while loading:" + c, e);
                     }
-                    List<Class<AssignmentEndpoint>> assignments = plugin.getAssignments();
+                    List<Class<AssignmentEndpoint>> assignments = plugin.getAssignments(c);
                     lesson.setAssignments(createAssignment(assignments));
                     lessons.add(lesson);
                     pluginEndpointPublisher.publish(plugin.getEndpoints());