Skip to content

Commit f1557ed

Browse files
impolitepandaimpolitepanda
authored
[backend] Added bulk create expectation traces endpoint (#2873)
1 parent 590a994 commit f1557ed

File tree

15 files changed

+661
-83
lines changed

15 files changed

+661
-83
lines changed

‎openbas-api/src/main/java/io/openbas/migration/V3_82__Add_Unique_constraint_injects_expectations_traces.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package io.openbas.migration;
2+
3+
import java.sql.Connection;
4+
import java.sql.Statement;
5+
import org.flywaydb.core.api.migration.BaseJavaMigration;
6+
import org.flywaydb.core.api.migration.Context;
7+
import org.springframework.stereotype.Component;
8+
9+
@Component
10+
public class V3_82__Add_Unique_constraint_injects_expectations_traces extends BaseJavaMigration {
11+
12+
@Override
13+
public void migrate(Context context) throws Exception {
14+
Connection connection = context.getConnection();
15+
Statement select = connection.createStatement();
16+
17+
select.execute(
18+
" DELETE FROM injects_expectations_traces iet WHERE exists ("
19+
+ " SELECT 1 from injects_expectations_traces iet1"
20+
+ " WHERE iet.inject_expectation_trace_expectation = iet1.inject_expectation_trace_expectation"
21+
+ " AND iet.inject_expectation_trace_source_id = iet1.inject_expectation_trace_source_id"
22+
+ " AND iet.inject_expectation_trace_alert_link = iet1.inject_expectation_trace_alert_link"
23+
+ " AND iet.inject_expectation_trace_alert_name = iet1.inject_expectation_trace_alert_name"
24+
+ ");");
25+
26+
select.execute(
27+
"ALTER TABLE injects_expectations_traces ADD CONSTRAINT unique_injects_expectations_traces_constraint UNIQUE (inject_expectation_trace_expectation, inject_expectation_trace_source_id, inject_expectation_trace_alert_link, inject_expectation_trace_alert_name);");
28+
}
29+
}

‎openbas-api/src/main/java/io/openbas/rest/expectation/ExpectationApi.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import io.openbas.database.model.InjectExpectation;
44
import io.openbas.rest.exercise.form.ExpectationUpdateInput;
55
import io.openbas.rest.helper.RestBehavior;
6+
import io.openbas.rest.inject.form.InjectExpectationBulkUpdateInput;
67
import io.openbas.rest.inject.form.InjectExpectationUpdateInput;
78
import io.openbas.service.InjectExpectationService;
89
import io.swagger.v3.oas.annotations.Operation;
@@ -139,4 +140,14 @@ public InjectExpectation updateInjectExpectation(
139140
@Valid @RequestBody @NotNull InjectExpectationUpdateInput input) {
140141
return injectExpectationService.updateInjectExpectation(expectationId, input);
141142
}
143+
144+
@Operation(
145+
summary = "Bulk Update Inject Expectation",
146+
description = "Bulk Update Inject expectation from an external source, e.g., EDR collector.")
147+
@PutMapping(INJECTS_EXPECTATIONS_URI + "/bulk")
148+
@Transactional(rollbackOn = Exception.class)
149+
public void updateInjectExpectation(
150+
@Valid @RequestBody @NotNull InjectExpectationBulkUpdateInput inputs) {
151+
injectExpectationService.bulkUpdateInjectExpectation(inputs.getInputs());
152+
}
142153
}

‎openbas-api/src/main/java/io/openbas/rest/inject/form/InjectExpectationBulkUpdateInput.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package io.openbas.rest.inject.form;
2+
3+
import com.fasterxml.jackson.annotation.JsonProperty;
4+
import jakarta.validation.constraints.NotNull;
5+
import java.util.Map;
6+
import lombok.AllArgsConstructor;
7+
import lombok.Builder;
8+
import lombok.Data;
9+
import lombok.NoArgsConstructor;
10+
11+
@AllArgsConstructor
12+
@NoArgsConstructor
13+
@Builder
14+
@Data
15+
public class InjectExpectationBulkUpdateInput {
16+
17+
/** Map of expectation IDs to their corresponding update inputs. */
18+
@NotNull
19+
@JsonProperty("inputs")
20+
private Map<String, InjectExpectationUpdateInput> inputs;
21+
}

‎openbas-api/src/main/java/io/openbas/rest/inject_expectation_trace/InjectExpectationTraceApi.java‎

Lines changed: 55 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,87 @@
11
package io.openbas.rest.inject_expectation_trace;
22

3+
import io.openbas.aop.LogExecutionTime;
34
import io.openbas.database.model.Collector;
45
import io.openbas.database.model.InjectExpectationTrace;
56
import io.openbas.database.repository.CollectorRepository;
6-
import io.openbas.database.repository.InjectExpectationRepository;
7+
import io.openbas.database.repository.InjectExpectationTraceRepository;
78
import io.openbas.rest.exception.ElementNotFoundException;
89
import io.openbas.rest.helper.RestBehavior;
10+
import io.openbas.rest.inject_expectation_trace.form.InjectExpectationTraceBulkInsertInput;
911
import io.openbas.rest.inject_expectation_trace.form.InjectExpectationTraceInput;
1012
import io.openbas.service.InjectExpectationTraceService;
13+
import io.swagger.v3.oas.annotations.Operation;
14+
import io.swagger.v3.oas.annotations.responses.ApiResponse;
15+
import io.swagger.v3.oas.annotations.responses.ApiResponses;
1116
import jakarta.validation.Valid;
17+
import jakarta.validation.constraints.NotNull;
1218
import java.util.List;
1319
import lombok.RequiredArgsConstructor;
20+
import lombok.extern.java.Log;
1421
import org.springframework.security.access.prepost.PreAuthorize;
1522
import org.springframework.web.bind.annotation.*;
1623

1724
@RequiredArgsConstructor
1825
@RestController
19-
@RequestMapping("/api/inject-expectations-traces")
26+
@RequestMapping(InjectExpectationTraceApi.INJECT_EXPECTATION_TRACES_URI)
2027
@PreAuthorize("isAdmin()")
28+
@Log
2129
public class InjectExpectationTraceApi extends RestBehavior {
2230

2331
public static final String INJECT_EXPECTATION_TRACES_URI = "/api/inject-expectations-traces";
2432

2533
private final InjectExpectationTraceService injectExpectationTraceService;
26-
private final InjectExpectationRepository injectExpectationRepository;
34+
private final InjectExpectationTraceRepository injectExpectationTraceRepository;
2735
private final CollectorRepository collectorRepository;
2836

37+
/**
38+
* @deprecated since 1.16.0, forRemoval = true
39+
* @see #bulkInsertInjectExpectationTraceForCollector(InjectExpectationTraceBulkInsertInput)
40+
*/
41+
@Deprecated(since = "1.16.0", forRemoval = true)
42+
@Operation(
43+
summary =
44+
"Create inject expectation trace for collector. Deprecated since 1.16.0. Replaced by "
45+
+ INJECT_EXPECTATION_TRACES_URI
46+
+ "/bulk")
2947
@PostMapping()
3048
public InjectExpectationTrace createInjectExpectationTraceForCollector(
3149
@Valid @RequestBody InjectExpectationTraceInput input) {
32-
InjectExpectationTrace injectExpectationTrace = new InjectExpectationTrace();
33-
injectExpectationTrace.setUpdateAttributes(input);
34-
injectExpectationTrace.setInjectExpectation(
35-
injectExpectationRepository
36-
.findById(input.getInjectExpectationId())
37-
.orElseThrow(() -> new ElementNotFoundException("Inject expectation not found")));
38-
Collector collector =
39-
collectorRepository
40-
.findById(input.getSourceId())
41-
.orElseThrow(() -> new ElementNotFoundException("Collector not found"));
42-
injectExpectationTrace.setSecurityPlatform(collector.getSecurityPlatform());
43-
return this.injectExpectationTraceService.createInjectExpectationTrace(injectExpectationTrace);
50+
51+
InjectExpectationTraceBulkInsertInput bulkInput = new InjectExpectationTraceBulkInsertInput();
52+
bulkInput.setExpectationTraces(List.of(input));
53+
54+
this.bulkInsertInjectExpectationTraceForCollector(bulkInput);
55+
// fetch the inserted data from the DB
56+
return this.injectExpectationTraceRepository
57+
.findByAlertLink(input.getAlertLink())
58+
.orElseThrow(ElementNotFoundException::new);
59+
}
60+
61+
/**
62+
* Bulk insert inject expectation traces for a collector.
63+
*
64+
* @param inputs the list of inject expectation trace inputs to be inserted
65+
*/
66+
@Operation(summary = "Bulk insert inject expectation traces")
67+
@ApiResponses(
68+
value = {
69+
@ApiResponse(
70+
responseCode = "200",
71+
description = "Inject expectation traces inserted successfully")
72+
})
73+
@LogExecutionTime
74+
@PostMapping("/bulk")
75+
public void bulkInsertInjectExpectationTraceForCollector(
76+
@Valid @RequestBody @NotNull InjectExpectationTraceBulkInsertInput inputs) {
77+
if (inputs.getExpectationTraces().isEmpty()) {
78+
return;
79+
}
80+
this.injectExpectationTraceService.bulkInsertInjectExpectationTraces(
81+
inputs.getExpectationTraces());
4482
}
4583

84+
@Operation(summary = "Get inject expectation traces from collector")
4685
@GetMapping()
4786
public List<InjectExpectationTrace> getInjectExpectationTracesFromCollector(
4887
@RequestParam String injectExpectationId, @RequestParam String sourceId) {
@@ -54,6 +93,7 @@ public List<InjectExpectationTrace> getInjectExpectationTracesFromCollector(
5493
injectExpectationId, collector.getSecurityPlatform().getId());
5594
}
5695

96+
@Operation(summary = "Get inject expectation traces' count")
5797
@GetMapping("/count")
5898
public long getAlertLinksNumber(
5999
@RequestParam String injectExpectationId,

‎openbas-api/src/main/java/io/openbas/rest/inject_expectation_trace/form/InjectExpectationTraceBulkInsertInput.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package io.openbas.rest.inject_expectation_trace.form;
2+
3+
import com.fasterxml.jackson.annotation.JsonProperty;
4+
import io.openbas.config.AppConfig;
5+
import jakarta.validation.constraints.NotNull;
6+
import java.util.List;
7+
import lombok.Data;
8+
9+
@Data
10+
public class InjectExpectationTraceBulkInsertInput {
11+
12+
@JsonProperty("expectation_traces")
13+
@NotNull(message = AppConfig.MANDATORY_MESSAGE)
14+
private List<InjectExpectationTraceInput> expectationTraces;
15+
}

‎openbas-api/src/main/java/io/openbas/service/InjectExpectationService.java‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,12 @@
2929
import java.util.stream.Collectors;
3030
import java.util.stream.Stream;
3131
import lombok.RequiredArgsConstructor;
32+
import lombok.extern.java.Log;
3233
import org.springframework.data.jpa.domain.Specification;
3334
import org.springframework.stereotype.Service;
3435
import org.springframework.transaction.annotation.Transactional;
3536

37+
@Log
3638
@RequiredArgsConstructor
3739
@Service
3840
public class InjectExpectationService {
@@ -418,6 +420,56 @@ public InjectExpectation updateInjectExpectation(
418420
return injectExpectation;
419421
}
420422

423+
public void bulkUpdateInjectExpectation(
424+
@Valid @NotNull Map<String, InjectExpectationUpdateInput> inputs) {
425+
if (inputs.isEmpty()) {
426+
return;
427+
}
428+
429+
List<InjectExpectation> injectExpectations =
430+
fromIterable(this.injectExpectationRepository.findAllById(inputs.keySet()));
431+
Map<String, InjectExpectation> expectationsToUpdate =
432+
injectExpectations.stream().collect(Collectors.toMap(InjectExpectation::getId, e -> e));
433+
434+
Collector collector =
435+
this.collectorRepository
436+
.findById(
437+
inputs.values().stream()
438+
.findFirst()
439+
.orElseThrow(ElementNotFoundException::new)
440+
.getCollectorId())
441+
.orElseThrow(ElementNotFoundException::new);
442+
443+
// Update inject expectation at agent level
444+
for (Map.Entry<String, InjectExpectationUpdateInput> entry : inputs.entrySet()) {
445+
String injectExpectationId = entry.getKey();
446+
InjectExpectationUpdateInput input = entry.getValue();
447+
448+
InjectExpectation injectExpectation = expectationsToUpdate.get(injectExpectationId);
449+
if (injectExpectation == null) {
450+
log.severe("Inject expectation not found for ID: " + injectExpectationId);
451+
continue;
452+
}
453+
454+
injectExpectation =
455+
this.computeExpectation(
456+
injectExpectation,
457+
collector.getId(),
458+
COLLECTOR,
459+
collector.getName(),
460+
input.getResult(),
461+
input.getIsSuccess(),
462+
input.getMetadata());
463+
464+
Inject inject = injectExpectation.getInject();
465+
// Compute potential expectations for asset
466+
propagateUpdateToAssets(injectExpectation, inject, collector);
467+
// Compute potential expectations for asset groups
468+
propagateUpdateToAssetGroups(inject, collector);
469+
// end of computing
470+
}
471+
}
472+
421473
private void propagateUpdateToAssets(
422474
InjectExpectation injectExpectation, Inject inject, Collector collector) {
423475
InjectExpectation finalInjectExpectation = injectExpectation;

‎openbas-api/src/main/java/io/openbas/service/InjectExpectationTraceService.java‎

Lines changed: 53 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,21 @@
11
package io.openbas.service;
22

3+
import io.openbas.database.model.Collector;
4+
import io.openbas.database.model.InjectExpectation;
35
import io.openbas.database.model.InjectExpectationTrace;
46
import io.openbas.database.model.SecurityPlatform;
7+
import io.openbas.database.raw.impl.SimpleRawExpectationTrace;
8+
import io.openbas.database.repository.CollectorRepository;
59
import io.openbas.database.repository.InjectExpectationTraceRepository;
610
import io.openbas.database.repository.SecurityPlatformRepository;
711
import io.openbas.rest.exception.ElementNotFoundException;
12+
import io.openbas.rest.inject_expectation_trace.form.InjectExpectationTraceInput;
813
import jakarta.validation.constraints.NotNull;
9-
import java.util.List;
10-
import java.util.Optional;
14+
import java.util.*;
1115
import lombok.RequiredArgsConstructor;
1216
import lombok.extern.java.Log;
1317
import org.springframework.stereotype.Service;
18+
import org.springframework.transaction.annotation.Transactional;
1419

1520
@Service
1621
@Log
@@ -19,24 +24,9 @@ public class InjectExpectationTraceService {
1924

2025
private final InjectExpectationTraceRepository injectExpectationTraceRepository;
2126
private final SecurityPlatformRepository securityPlatformRepository;
22-
private static final String COLLECTOR_TYPE = "collector";
27+
private final CollectorRepository collectorRepository;
2328

24-
public InjectExpectationTrace createInjectExpectationTrace(
25-
@NotNull InjectExpectationTrace injectExpectationTrace) {
26-
Optional<InjectExpectationTrace> existingTrace =
27-
this.injectExpectationTraceRepository
28-
.findByAlertLinkAndAlertNameAndSecurityPlatformAndInjectExpectation(
29-
injectExpectationTrace.getAlertLink(),
30-
injectExpectationTrace.getAlertName(),
31-
injectExpectationTrace.getSecurityPlatform(),
32-
injectExpectationTrace.getInjectExpectation());
33-
if (existingTrace.isPresent()) {
34-
log.info("Existing trace present, no creation");
35-
return existingTrace.get();
36-
} else {
37-
return this.injectExpectationTraceRepository.save(injectExpectationTrace);
38-
}
39-
}
29+
private static final String COLLECTOR_TYPE = "collector";
4030

4131
public List<InjectExpectationTrace> getInjectExpectationTracesFromCollector(
4232
@NotNull String injectExpectationId, @NotNull String sourceId) {
@@ -59,4 +49,48 @@ public long getAlertLinksNumber(
5949
return this.injectExpectationTraceRepository.countAlerts(injectExpectationId, sourceId);
6050
}
6151
}
52+
53+
@Transactional(rollbackFor = Exception.class)
54+
public void bulkInsertInjectExpectationTraces(
55+
@NotNull List<InjectExpectationTraceInput> injectExpectationTraces) {
56+
if (injectExpectationTraces.isEmpty()) {
57+
return;
58+
}
59+
// We start by deduplicating the data, to avoid duplicates in the database
60+
// Convert the input list to InjectExpectationTrace objects and extract oldest trace's date
61+
// Start by getting the collector. We can take the first one since they are all the same
62+
Collector collector =
63+
collectorRepository
64+
.findById(injectExpectationTraces.getFirst().getSourceId())
65+
.orElseThrow(() -> new ElementNotFoundException("Collector not found"));
66+
Map<SimpleRawExpectationTrace, InjectExpectationTrace> traces = new HashMap<>();
67+
injectExpectationTraces.forEach(
68+
input -> {
69+
// Convert input to InjectExpectationTrace
70+
InjectExpectationTrace trace = new InjectExpectationTrace();
71+
trace.setUpdateAttributes(input);
72+
trace.setSecurityPlatform(collector.getSecurityPlatform());
73+
// We don't need to fetch the actual expectation here, we can just set the id as there is
74+
// no cascade
75+
trace.setInjectExpectation(new InjectExpectation());
76+
trace.getInjectExpectation().setId(input.getInjectExpectationId());
77+
78+
SimpleRawExpectationTrace simpleTrace = SimpleRawExpectationTrace.of(trace);
79+
80+
traces.computeIfAbsent(simpleTrace, k -> trace);
81+
});
82+
83+
// Save the remaining traces
84+
for (InjectExpectationTrace trace : traces.values()) {
85+
this.injectExpectationTraceRepository.insertIfNotExists(
86+
UUID.randomUUID().toString(),
87+
trace.getInjectExpectation().getId(),
88+
trace.getSecurityPlatform().getId(),
89+
trace.getAlertLink(),
90+
trace.getAlertName(),
91+
trace.getAlertDate(),
92+
trace.getCreatedAt(),
93+
trace.getUpdatedAt());
94+
}
95+
}
6296
}

0 commit comments

Comments
 (0)