[backend] remove duplicate code

Signed-off-by: Marine LM <[email protected]>
OpenBAS-Platform · Jan 30, 2025 · ea710d4 · ea710d4
1 parent 72e1416
commit ea710d4
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 78 deletions.
diff --git a/openbas-api/src/main/java/io/openbas/injectors/caldera/CalderaExecutor.java b/openbas-api/src/main/java/io/openbas/injectors/caldera/CalderaExecutor.java
@@ -16,7 +16,6 @@
 import io.openbas.database.model.*;
 import io.openbas.database.model.InjectExpectation.EXPECTATION_TYPE;
 import io.openbas.database.model.PayloadCommandBlock;
-import io.openbas.database.repository.InjectRepository;
 import io.openbas.execution.ExecutableInject;
 import io.openbas.executors.Injector;
 import io.openbas.injectors.caldera.client.model.Ability;
@@ -29,6 +28,7 @@
 import io.openbas.model.expectation.DetectionExpectation;
 import io.openbas.model.expectation.ManualExpectation;
 import io.openbas.model.expectation.PreventionExpectation;
+import io.openbas.rest.inject.service.InjectService;
 import io.openbas.service.AssetGroupService;
 import io.openbas.service.EndpointService;
 import io.openbas.service.InjectExpectationService;
@@ -55,7 +55,7 @@ public class CalderaExecutor extends Injector {
   private final EndpointService endpointService;
   private final AssetGroupService assetGroupService;
   private final InjectExpectationService injectExpectationService;
-  private final InjectRepository injectRepository;
+  private final InjectService injectService;
 
   @Override
   @Transactional
@@ -67,10 +67,9 @@ public ExecutionProcess process(
         content.getObfuscator() != null
             ? content.getObfuscator()
             : CalderaInjectContent.getDefaultObfuscator();
-    Inject inject =
-        this.injectRepository.findById(injection.getInjection().getInject().getId()).orElseThrow();
+    Inject inject = this.injectService.inject(injection.getInjection().getInject().getId());
 
-    Map<Endpoint, Boolean> assets = this.resolveAllAssets(injection);
+    Map<Endpoint, Boolean> assets = this.injectService.resolveAllAssetsToExecute(inject);
     // Execute inject for all assets
     if (assets.isEmpty()) {
       execution.addTrace(
@@ -151,7 +150,7 @@ public ExecutionProcess process(
                                 getNewInfoTrace(
                                     exploitResult.getCommand(),
                                     ExecutionTraceAction.EXECUTION,
-                                    asset.getAgents().getFirst(),
+                                    ((Endpoint) asset).getAgents().getFirst(),
                                     List.of()));
                             // Compute expectations
                             boolean isInGroup =
@@ -343,29 +342,6 @@ public StatusPayload getPayloadOutput(String externalId) {
 
   // -- PRIVATE --
 
-  private Map<Endpoint, Boolean> resolveAllAssets(@NotNull final ExecutableInject inject) {
-    Map<Endpoint, Boolean> assets = new HashMap<>();
-    inject
-        .getAssets()
-        .forEach(
-            (asset -> {
-              assets.put((Endpoint) Hibernate.unproxy(asset), false);
-            }));
-    inject
-        .getAssetGroups()
-        .forEach(
-            (assetGroup -> {
-              List<Asset> assetsFromGroup =
-                  this.assetGroupService.assetsFromAssetGroup(assetGroup.getId());
-              // Verify asset validity
-              assetsFromGroup.forEach(
-                  (asset) -> {
-                    assets.put((Endpoint) Hibernate.unproxy(asset), true);
-                  });
-            }));
-    return assets;
-  }
-
   private Endpoint findAndRegisterAssetForExecution(
       @NotNull final Inject inject, @NotNull final Asset asset) throws InterruptedException {
     Endpoint endpointForExecution = null;

diff --git a/openbas-api/src/main/java/io/openbas/injectors/openbas/OpenBASImplantExecutor.java b/openbas-api/src/main/java/io/openbas/injectors/openbas/OpenBASImplantExecutor.java
@@ -19,6 +19,7 @@
 import io.openbas.model.expectation.DetectionExpectation;
 import io.openbas.model.expectation.ManualExpectation;
 import io.openbas.model.expectation.PreventionExpectation;
+import io.openbas.rest.inject.service.InjectService;
 import io.openbas.service.AssetGroupService;
 import io.openbas.service.InjectExpectationService;
 import jakarta.validation.constraints.NotNull;
@@ -36,29 +37,7 @@ public class OpenBASImplantExecutor extends Injector {
   private final AssetGroupService assetGroupService;
   private final InjectRepository injectRepository;
   private final InjectExpectationService injectExpectationService;
-
-  private Map<Asset, Boolean> resolveAllAssets(@NotNull final ExecutableInject inject) {
-    Map<Asset, Boolean> assets = new HashMap<>();
-    inject
-        .getAssets()
-        .forEach(
-            (asset -> {
-              assets.put(asset, false);
-            }));
-    inject
-        .getAssetGroups()
-        .forEach(
-            (assetGroup -> {
-              List<Asset> assetsFromGroup =
-                  this.assetGroupService.assetsFromAssetGroup(assetGroup.getId());
-              // Verify asset validity
-              assetsFromGroup.forEach(
-                  (asset) -> {
-                    assets.put(asset, true);
-                  });
-            }));
-    return assets;
-  }
+  private final InjectService injectService;
 
   /** In case of direct asset, we have an individual expectation for the asset */
   private void computeExpectationsForAsset(
@@ -227,7 +206,7 @@ public ExecutionProcess process(Execution execution, ExecutableInject injection)
       throws Exception {
     Inject inject =
         this.injectRepository.findById(injection.getInjection().getInject().getId()).orElseThrow();
-    Map<Asset, Boolean> assets = this.resolveAllAssets(injection);
+    Map<Endpoint, Boolean> assets = this.injectService.resolveAllAssetsToExecute(inject);
 
     // Check assets target
     if (assets.isEmpty()) {

diff --git a/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectService.java b/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectService.java
@@ -34,10 +34,7 @@
 import jakarta.validation.constraints.NotBlank;
 import jakarta.validation.constraints.NotNull;
 import java.time.Instant;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -93,6 +90,29 @@ public void deleteAll(List<Inject> injects) {
     }
   }
 
+  public Map<Endpoint, Boolean> resolveAllAssetsToExecute(@NotNull final Inject inject) {
+    Map<Endpoint, Boolean> assets = new HashMap<>();
+    inject
+        .getAssets()
+        .forEach(
+            (asset -> {
+              assets.put((Endpoint) asset, false);
+            }));
+    inject
+        .getAssetGroups()
+        .forEach(
+            (assetGroup -> {
+              List<Asset> assetsFromGroup =
+                  this.assetGroupService.assetsFromAssetGroup(assetGroup.getId());
+              // Verify asset validity
+              assetsFromGroup.forEach(
+                  (asset) -> {
+                    assets.put((Endpoint) asset, true);
+                  });
+            }));
+    return assets;
+  }
+
   public void cleanInjectsDocExercise(String exerciseId, String documentId) {
     // Delete document from all exercise injects
     List<Inject> exerciseInjects =

diff --git a/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectStatusService.java b/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectStatusService.java
@@ -7,6 +7,7 @@
 import io.openbas.database.repository.AgentRepository;
 import io.openbas.database.repository.InjectRepository;
 import io.openbas.database.repository.InjectStatusRepository;
+import io.openbas.execution.ExecutableInject;
 import io.openbas.rest.exception.ElementNotFoundException;
 import io.openbas.rest.inject.form.InjectExecutionAction;
 import io.openbas.rest.inject.form.InjectExecutionInput;
@@ -27,6 +28,7 @@ public class InjectStatusService {
   private final AgentRepository agentRepository;
   private final InjectUtils injectUtils;
   private final InjectStatusRepository injectStatusRepository;
+  private final InjectService injectService;
 
   public List<InjectStatus> findPendingInjectStatusByType(String injectType) {
     return this.injectStatusRepository.pendingForInjectType(injectType);
@@ -74,25 +76,10 @@ private int getCompleteTrace(Inject inject) {
         .size();
   }
 
-  private int calculateTotalAssets(Inject inject) {
-    int assetCount = Optional.ofNullable(inject.getAssets()).map(List::size).orElse(0);
-    int groupAssetCount =
-        Optional.ofNullable(inject.getAssetGroups())
-            .map(
-                groups ->
-                    groups.stream()
-                        .mapToInt(
-                            group ->
-                                Optional.ofNullable(group.getAssets()).map(List::size).orElse(0))
-                        .sum())
-            .orElse(0);
-    return assetCount + groupAssetCount;
-  }
-
   public boolean isAllInjectAssetsExecuted(Inject inject) {
     int totalCompleteTrace = getCompleteTrace(inject);
-    int totalAssetCount = calculateTotalAssets(inject);
-    return totalAssetCount == totalCompleteTrace;
+    Map<Endpoint, Boolean> assets = this.injectService.resolveAllAssetsToExecute(inject);
+    return assets.size() == totalCompleteTrace;
   }
 
   public void updateFinalInjectStatus(InjectStatus injectStatus, Instant finishTime) {
@@ -226,6 +213,7 @@ public InjectStatus initializeInjectStatus(
     InjectStatus injectStatus = getOrInitializeInjectStatus(inject);
 
     if (trace != null) {
+      trace.setInjectStatus(injectStatus);
       injectStatus.addTrace(trace);
     }
     injectStatus.setName(status);

diff --git a/openbas-api/src/main/java/io/openbas/scheduler/jobs/InjectsExecutionJob.java b/openbas-api/src/main/java/io/openbas/scheduler/jobs/InjectsExecutionJob.java
@@ -181,11 +181,12 @@ private void setInjectStatusAndExecuteInject(ExecutableInject executableInject,
               LOGGER.log(Level.INFO, "Executing inject " + inject.getInject().getTitle());
               // Executor logics
               ExecutableInject newExecutableInject = executableInject;
+
+              InjectStatus injectStatus =
+                  injectStatusService.initializeInjectStatus(
+                      inject.getId(), ExecutionStatus.EXECUTING, null);
               if (Boolean.TRUE.equals(injectorContract.getNeedsExecutor())) {
                 // Status
-                InjectStatus injectStatus =
-                    injectStatusService.initializeInjectStatus(
-                        inject.getId(), ExecutionStatus.EXECUTING, null);
                 inject.setStatus(injectStatus);
                 try {
                   newExecutableInject =

diff --git a/openbas-model/src/main/java/io/openbas/database/model/InjectStatus.java b/openbas-model/src/main/java/io/openbas/database/model/InjectStatus.java
@@ -52,7 +52,6 @@ public Map<String, Agent> getStatusMapIdentifierAgent() {
 
   public void addTrace(ExecutionTraces trace) {
     this.getTraces().add(trace);
-    trace.setInjectStatus(this);
   }
 
   public void addTrace(