[backend] Fix MITRE ATT&CK matrix dashboard results (#1670)

OpenBAS-Platform · Oct 17, 2024 · 033f92b · 033f92b
1 parent 28f035a
commit 033f92b
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 42 deletions.
diff --git a/openbas-api/src/main/java/io/openbas/rest/statistic/StatisticApi.java b/openbas-api/src/main/java/io/openbas/rest/statistic/StatisticApi.java
@@ -26,11 +26,13 @@
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 import static io.openbas.config.SessionHelper.currentUser;
 import static io.openbas.helper.StreamHelper.fromIterable;
+import static io.openbas.utils.AtomicTestingUtils.getExpectationResultByTypesFromRaw;
+import static java.util.stream.Collectors.groupingBy;
 
 @RestController
 @RequiredArgsConstructor
@@ -96,13 +98,13 @@ private StatisticElement computeUserStat(Instant from, StatisticRepository repos
 
   private List<ExpectationResultsByType> computeGlobalExpectationResults(@NotNull final Instant from) {
     List<RawInjectExpectation> rawInjectExpectations = fromIterable(this.exerciseRepository.allInjectExpectationsFromDate(from));
-    return AtomicTestingUtils.getExpectationResultByTypesFromRaw(rawInjectExpectations);
+    return getExpectationResultByTypesFromRaw(rawInjectExpectations);
   }
 
   private List<ExpectationResultsByType> computeUserExpectationResults(@NotNull final Instant from) {
     OpenBASPrincipal user = currentUser();
     List<RawInjectExpectation> rawInjectExpectations = fromIterable(this.exerciseRepository.allGrantedInjectExpectationsFromDate(from, user.getId()));
-    return AtomicTestingUtils.getExpectationResultByTypesFromRaw(rawInjectExpectations);
+    return getExpectationResultByTypesFromRaw(rawInjectExpectations);
   }
 
   private List<InjectExpectationResultsByAttackPattern> computeGlobalInjectExpectationResults(
@@ -118,43 +120,58 @@ private List<InjectExpectationResultsByAttackPattern> computeUserInjectExpectati
     return injectExpectationResultsByAttackPatternFromRawGlobalInjectExpectation(rawGlobalInjectExpectations);
   }
 
-  private List<InjectExpectationResultsByAttackPattern> injectExpectationResultsByAttackPatternFromRawGlobalInjectExpectation(List<RawGlobalInjectExpectation> rawGlobalInjectExpectations ) {
+  private List<InjectExpectationResultsByAttackPattern> injectExpectationResultsByAttackPatternFromRawGlobalInjectExpectation(
+      List<RawGlobalInjectExpectation> rawGlobalInjectExpectations) {
     return rawGlobalInjectExpectations.stream()
-            .map(RawGlobalInjectExpectation::getAttack_pattern_id)
-            .distinct()
-            .map(
-                    attackPatternId ->
-                    {
-                      InjectExpectationResultsByAttackPattern resultExpectation = new InjectExpectationResultsByAttackPattern();
-                      resultExpectation.setAttackPattern(new AttackPattern());
-                      resultExpectation.getAttackPattern().setId(attackPatternId);
-                      resultExpectation.setResults(rawGlobalInjectExpectations.stream()
-                              .filter((expectation) -> expectation.getAttack_pattern_id().equals(attackPatternId))
-                              .map((expectation) -> {
-                                InjectExpectationResultsByAttackPattern.InjectExpectationResultsByType resultInjectExpectationResultsByAttackPattern = new InjectExpectationResultsByAttackPattern.InjectExpectationResultsByType();
-                                resultInjectExpectationResultsByAttackPattern.setInjectTitle(expectation.getInject_title());
-                                if(expectation.getInject_expectation_type() != null) {
-                                  SimpleRawInjectExpectation rawInjectExpectation = new SimpleRawInjectExpectation();
-                                  rawInjectExpectation.setInject_expectation_score(expectation.getInject_expectation_score());
-                                  rawInjectExpectation.setInject_expectation_expected_score(expectation.getInject_expectation_expected_score());
-                                  rawInjectExpectation.setInject_expectation_type(expectation.getInject_expectation_type());
-                                  resultInjectExpectationResultsByAttackPattern
-                                          .setResults(AtomicTestingUtils.getExpectationResultByTypesFromRaw(Stream.of(rawInjectExpectation)
-                                                  .collect(Collectors.toList())));
-                                } else {
-                                  resultInjectExpectationResultsByAttackPattern
-                                          .setResults(AtomicTestingUtils.getExpectationResultByTypesFromRaw(new ArrayList<>()));
-                                }
-                                return resultInjectExpectationResultsByAttackPattern;
-                              })
-                              .collect(Collectors.toList())
-                      );
-
-                      return resultExpectation;
-                    }
-            )
-            .collect(Collectors.toList())
-            ;
+        .map(RawGlobalInjectExpectation::getAttack_pattern_id)
+        .distinct()
+        .map(
+            attackPatternId ->
+            {
+              InjectExpectationResultsByAttackPattern resultExpectation = new InjectExpectationResultsByAttackPattern();
+              resultExpectation.setAttackPattern(new AttackPattern());
+              resultExpectation.getAttackPattern().setId(attackPatternId);
+
+              Map<String, Map<String, List<RawGlobalInjectExpectation>>> rawGlobalInjectExpectationsGroupByAttackAndInjectId = rawGlobalInjectExpectations.stream()
+                  .collect(groupingBy(RawGlobalInjectExpectation::getAttack_pattern_id,
+                      groupingBy(RawGlobalInjectExpectation::getInject_id)));
+
+              List<InjectExpectationResultsByAttackPattern.InjectExpectationResultsByType> results = new ArrayList<>();
+
+              rawGlobalInjectExpectationsGroupByAttackAndInjectId.forEach((attackId, injects) -> {
+                if (attackId.equals(attackPatternId)) {
+                  injects.forEach((injectId, expectations) -> {
+                    RawGlobalInjectExpectation expectation = expectations.getFirst();
+                    InjectExpectationResultsByAttackPattern.InjectExpectationResultsByType resultInjectExpectationResultsByAttackPattern = new InjectExpectationResultsByAttackPattern.InjectExpectationResultsByType();
+                    resultInjectExpectationResultsByAttackPattern.setInjectTitle(expectation.getInject_title());
+
+                    ArrayList<RawInjectExpectation> expectationsRefined = new ArrayList<>();
+
+                    expectations.stream().forEach(e -> {
+                      if (e.getInject_expectation_type() != null) {
+                        SimpleRawInjectExpectation rawInjectExpectation = new SimpleRawInjectExpectation();
+                        rawInjectExpectation.setInject_expectation_score(e.getInject_expectation_score());
+                        rawInjectExpectation.setInject_expectation_expected_score(
+                            e.getInject_expectation_expected_score());
+                        rawInjectExpectation.setInject_expectation_type(e.getInject_expectation_type());
+                        expectationsRefined.add(rawInjectExpectation);
+                      }
+                    });
+
+                    resultInjectExpectationResultsByAttackPattern
+                        .setResults(AtomicTestingUtils.getExpectationResultByTypesFromRaw(expectationsRefined));
+
+                    results.add(resultInjectExpectationResultsByAttackPattern);
+                  });
+                }
+              });
+
+              resultExpectation.setResults(results);
+
+              return resultExpectation;
+            }
+        )
+        .collect(Collectors.toList());
   }
 
 }
diff --git a/openbas-model/src/main/java/io/openbas/database/raw/RawGlobalInjectExpectation.java b/openbas-model/src/main/java/io/openbas/database/raw/RawGlobalInjectExpectation.java
@@ -8,6 +8,8 @@ public interface RawGlobalInjectExpectation {
 
     Double getInject_expectation_expected_score();
 
+    String getInject_id();
+
     String getInject_title();
 
     String getAttack_pattern_id();

diff --git a/openbas-model/src/main/java/io/openbas/database/repository/ExerciseRepository.java b/openbas-model/src/main/java/io/openbas/database/repository/ExerciseRepository.java
@@ -96,11 +96,12 @@ List<RawInjectExpectation> allGrantedInjectExpectationsFromDate(@Param("from") I
    * @return a list of expectations
    */
   @Query(value =
-      "SELECT ie.inject_expectation_type, ie.inject_expectation_score, ie.inject_expectation_expected_score, injects.inject_title, icap.attack_pattern_id "
+      "SELECT ie.inject_expectation_type, ie.inject_expectation_score, ie.inject_expectation_expected_score, "
+          + "injects.inject_id, injects.inject_title, icap.attack_pattern_id "
           +
           "FROM exercises " +
           "INNER JOIN injects ON exercises.exercise_id = injects.inject_exercise " +
-          "LEFT JOIN injects_expectations ie ON injects.inject_id = ie.exercise_id " +
+          "JOIN injects_expectations ie ON injects.inject_id = ie.inject_id " +
           "INNER JOIN injectors_contracts ic ON injects.inject_injector_contract = ic.injector_contract_id " +
           "INNER JOIN injectors_contracts_attack_patterns icap ON ic.injector_contract_id = icap.injector_contract_id "
           +
@@ -115,7 +116,8 @@ List<RawInjectExpectation> allGrantedInjectExpectationsFromDate(@Param("from") I
    * @return the list of global expectations
    */
   @Query(value =
-      "SELECT ie.inject_expectation_type, ie.inject_expectation_score, ie.inject_expectation_expected_score, injects.inject_title, icap.attack_pattern_id "
+      "SELECT ie.inject_expectation_type, ie.inject_expectation_score, ie.inject_expectation_expected_score, "
+          + "injects.inject_id, injects.inject_title, icap.attack_pattern_id "
           +
           "FROM exercises " +
           "INNER JOIN injects ON exercises.exercise_id = injects.inject_exercise " +