Repository: OpenAEV-Platform/openaev (java)

Issue:
Scenario by AI with Data Team
## Context
~~This collector will allow to sync between AI Model with python collector to OpenBAS Platform~~


Discussion:
Update: We'll use http call without injector to call the AI webservice.

Patch:
diff --git a/openbas-api/src/main/java/io/openbas/config/AppConfig.java b/openbas-api/src/main/java/io/openbas/config/AppConfig.java
--- a/openbas-api/src/main/java/io/openbas/config/AppConfig.java
+++ b/openbas-api/src/main/java/io/openbas/config/AppConfig.java
@@ -12,6 +12,7 @@
 import org.springframework.scheduling.annotation.EnableScheduling;
 import org.springframework.stereotype.Component;
 import org.springframework.transaction.annotation.EnableTransactionManagement;
+import org.springframework.web.client.RestTemplate;
 
 @Component
 @EnableAsync
@@ -49,4 +50,9 @@ public OpenAPI openBASOpenAPI() {
                 .description("OpenBAS documentation")
                 .url("https://docs.openbas.io/"));
   }
+
+  @Bean
+  public RestTemplate restTemplate() {
+    return new RestTemplate();
+  }
 }
diff --git a/openbas-api/src/main/java/io/openbas/ee/Ee.java b/openbas-api/src/main/java/io/openbas/ee/Ee.java
--- a/openbas-api/src/main/java/io/openbas/ee/Ee.java
+++ b/openbas-api/src/main/java/io/openbas/ee/Ee.java
@@ -156,20 +156,27 @@ public boolean isLicenseActive(License license) {
         && (Instant.now().isBefore(license.getExpirationDate()) || license.isExtraExpiration());
   }
 
+  public String getEnterpriseEditionLicensePem() {
+    Setting pemSetting =
+        this.settingRepository.findByKey(PLATFORM_ENTERPRISE_LICENSE.key()).orElse(new Setting());
+    String pem = pemSetting.getValue();
+    String pemFromConfig = openBASConfig.getApplicationLicense();
+    if (pemFromConfig != null && !pemFromConfig.trim().isEmpty()) {
+      return pemFromConfig.trim();
+    }
+    if (pem != null && !pem.trim().isEmpty()) {
+      return pem.trim();
+    }
+    return "";
+  }
+
   /**
    * This should not be called directly. Prefer to use the license Cache Manager:
    *
    * @see(LicenseCacheManager#getEnterpriseEditionInfo)
    */
   public License getEnterpriseEditionInfo() {
-    Map<String, Setting> dbSettings = mapOfSettings(fromIterable(this.settingRepository.findAll()));
-    String pem =
-        ofNullable(dbSettings.get(PLATFORM_ENTERPRISE_LICENSE.key()))
-            .map(Setting::getValue)
-            .orElse(PLATFORM_ENTERPRISE_LICENSE.defaultValue());
-    String pemFromConfig = openBASConfig.getApplicationLicense();
-    boolean isLicenseByConfig = pemFromConfig != null && !pemFromConfig.trim().isEmpty();
-    String certificatePem = isLicenseByConfig ? pemFromConfig.trim() : pem.trim();
+    String certificatePem = getEnterpriseEditionLicensePem();
     try {
       if (certificatePem.isEmpty()) {
         throw new IllegalArgumentException("Certificate Pem is null or empty");
@@ -177,7 +184,7 @@ public License getEnterpriseEditionInfo() {
       return getEnterpriseEditionInfoFromPem(certificatePem);
     } catch (Exception e) {
       License license = new License();
-      license.setLicenseByConfiguration(isLicenseByConfig);
+      license.setLicenseByConfiguration(false);
       return license;
     }
   }
diff --git a/openbas-api/src/main/java/io/openbas/helper/RabbitMQHelper.java b/openbas-api/src/main/java/io/openbas/helper/RabbitMQHelper.java
--- a/openbas-api/src/main/java/io/openbas/helper/RabbitMQHelper.java
+++ b/openbas-api/src/main/java/io/openbas/helper/RabbitMQHelper.java
@@ -55,7 +55,7 @@ public static String getRabbitMQVersion(RabbitmqConfig rabbitmqConfig) {
 
       RestTemplate restTemplate;
       try {
-        restTemplate = restTemplate(rabbitmqConfig);
+        restTemplate = rabbitMQRestTemplate(rabbitmqConfig);
       } catch (KeyStoreException
           | NoSuchAlgorithmException
           | KeyManagementException
@@ -89,7 +89,7 @@ public static String getRabbitMQVersion(RabbitmqConfig rabbitmqConfig) {
     return rabbitMQVersion;
   }
 
-  private static RestTemplate restTemplate(RabbitmqConfig rabbitmqConfig)
+  private static RestTemplate rabbitMQRestTemplate(RabbitmqConfig rabbitmqConfig)
       throws KeyStoreException,
           NoSuchAlgorithmException,
           KeyManagementException,
diff --git a/openbas-api/src/main/java/io/openbas/rest/attack_pattern/AttackPatternApi.java b/openbas-api/src/main/java/io/openbas/rest/attack_pattern/AttackPatternApi.java
--- a/openbas-api/src/main/java/io/openbas/rest/attack_pattern/AttackPatternApi.java
+++ b/openbas-api/src/main/java/io/openbas/rest/attack_pattern/AttackPatternApi.java
@@ -18,10 +18,13 @@
 import io.openbas.rest.attack_pattern.form.AttackPatternCreateInput;
 import io.openbas.rest.attack_pattern.form.AttackPatternUpdateInput;
 import io.openbas.rest.attack_pattern.form.AttackPatternUpsertInput;
+import io.openbas.rest.attack_pattern.service.AttackPatternService;
 import io.openbas.rest.exception.ElementNotFoundException;
 import io.openbas.rest.helper.RestBehavior;
 import io.openbas.utils.FilterUtilsJpa;
 import io.openbas.utils.pagination.SearchPaginationInput;
+import io.swagger.v3.oas.annotations.Operation;
+import jakarta.annotation.Nullable;
 import jakarta.transaction.Transactional;
 import jakarta.validation.Valid;
 import jakarta.validation.constraints.NotBlank;
@@ -29,39 +32,27 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
-import org.springframework.beans.factory.annotation.Autowired;
+import lombok.RequiredArgsConstructor;
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.jpa.domain.Specification;
 import org.springframework.security.access.annotation.Secured;
 import org.springframework.web.bind.annotation.*;
+import org.springframework.web.multipart.MultipartFile;
 
 @RestController
 @Secured(ROLE_USER)
+@RequiredArgsConstructor
 public class AttackPatternApi extends RestBehavior {
 
   public static final String ATTACK_PATTERN_URI = "/api/attack_patterns";
 
-  private AttackPatternRepository attackPatternRepository;
+  private final AttackPatternService attackPatternService;
 
-  private InjectorContractRepository injectorContractRepository;
-  private KillChainPhaseRepository killChainPhaseRepository;
-
-  @Autowired
-  public void setAttackPatternRepository(AttackPatternRepository attackPatternRepository) {
-    this.attackPatternRepository = attackPatternRepository;
-  }
-
-  @Autowired
-  public void setKillChainPhaseRepository(KillChainPhaseRepository killChainPhaseRepository) {
-    this.killChainPhaseRepository = killChainPhaseRepository;
-  }
-
-  @Autowired
-  public void setInjectorContractRepository(InjectorContractRepository injectorContractRepository) {
-    this.injectorContractRepository = injectorContractRepository;
-  }
+  private final AttackPatternRepository attackPatternRepository;
+  private final InjectorContractRepository injectorContractRepository;
+  private final KillChainPhaseRepository killChainPhaseRepository;
 
   @GetMapping("/api/attack_patterns")
   public List<RawAttackPattern> attackPatterns() {
@@ -78,6 +69,17 @@ public Page<AttackPattern> attackPatterns(
         AttackPattern.class);
   }
 
+  @PostMapping(ATTACK_PATTERN_URI + "/search-with-ai")
+  @Operation(
+      summary = "Extract Attack Paterns from text or files using AI",
+      description = "Get attack patterns ids extracted from a text or files using AI")
+  public List<String> searchAttackPatternWithTTPAIWebservice(
+      @RequestPart("files") @Nullable List<MultipartFile> files,
+      @RequestPart("text") @Nullable final String text) {
+    return attackPatternService.searchAttackPatternWithTTPAIWebservice(
+        files == null ? new ArrayList<>() : files, text == null ? "" : text);
+  }
+
   @GetMapping("/api/attack_patterns/{attackPatternId}")
   public AttackPattern attackPattern(@PathVariable String attackPatternId) {
     return attackPatternRepository
diff --git a/openbas-api/src/main/java/io/openbas/rest/attack_pattern/form/AnalysisResultFromTTPExtractionAIWebserviceOutput.java b/openbas-api/src/main/java/io/openbas/rest/attack_pattern/form/AnalysisResultFromTTPExtractionAIWebserviceOutput.java
new file mode 100644
--- /dev/null
+++ b/openbas-api/src/main/java/io/openbas/rest/attack_pattern/form/AnalysisResultFromTTPExtractionAIWebserviceOutput.java
@@ -0,0 +1,15 @@
+package io.openbas.rest.attack_pattern.form;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.Map;
+import lombok.Getter;
+
+// TODO ADD DOC
+@Getter
+public class AnalysisResultFromTTPExtractionAIWebserviceOutput {
+  @JsonProperty("text")
+  String text;
+
+  @JsonProperty("predictions")
+  Map<String, Number> predictions;
+}
diff --git a/openbas-api/src/main/java/io/openbas/rest/attack_pattern/service/AttackPatternService.java b/openbas-api/src/main/java/io/openbas/rest/attack_pattern/service/AttackPatternService.java
new file mode 100644
--- /dev/null
+++ b/openbas-api/src/main/java/io/openbas/rest/attack_pattern/service/AttackPatternService.java
@@ -0,0 +1,168 @@
+package io.openbas.rest.attack_pattern.service;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.openbas.database.model.AttackPattern;
+import io.openbas.database.repository.AttackPatternRepository;
+import io.openbas.ee.Ee;
+import io.openbas.rest.attack_pattern.form.AnalysisResultFromTTPExtractionAIWebserviceOutput;
+import jakarta.annotation.Resource;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.*;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.core.env.Environment;
+import org.springframework.core.io.ByteArrayResource;
+import org.springframework.http.HttpEntity;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
+import org.springframework.http.client.MultipartBodyBuilder;
+import org.springframework.stereotype.Service;
+import org.springframework.util.MultiValueMap;
+import org.springframework.web.client.RestClientException;
+import org.springframework.web.client.RestTemplate;
+import org.springframework.web.multipart.MultipartFile;
+
+@Slf4j
+@Service
+@RequiredArgsConstructor
+public class AttackPatternService {
+
+  @Resource protected ObjectMapper mapper;
+
+  private final Environment env;
+  private final AttackPatternRepository attackPatternRepository;
+  private final Ee ee;
+  private final RestTemplate restTemplate;
+
+  /**
+   * Call the TTP Extraction AI Webservice to analyze files and text input.
+   *
+   * @param files List of files to be analyzed, maximum 5 files.
+   * @param text Text input to be analyzed.
+   * @return Response body from the TTP Extraction AI Webservice, expected to be a JSON array
+   * @throws IOException
+   */
+  private String callTTPExtractionAIWebservice(List<MultipartFile> files, String text)
+      throws IOException {
+    String url = Objects.requireNonNull(env.getProperty("ttp.extraction.ai.webservice.url"));
+    String certificate = ee.getEnterpriseEditionLicensePem();
+    if (certificate == null || certificate.isBlank()) {
+      throw new IllegalStateException("Enterprise Edition is not available");
+    }
+    String encodedCertificate =
+        Base64.getEncoder().encodeToString(certificate.getBytes(StandardCharsets.UTF_8));
+
+    // Set up the headers for the request
+    HttpHeaders headers = new HttpHeaders();
+    headers.setContentType(MediaType.MULTIPART_FORM_DATA);
+    headers.add("X-OpenBAS-Certificate", encodedCertificate);
+
+    // Set up the request body
+    MultipartBodyBuilder bodyBuilder = new MultipartBodyBuilder();
+    for (MultipartFile file : files) {
+      ByteArrayResource resource =
+          new ByteArrayResource(file.getBytes()) {
+            @Override
+            public String getFilename() {
+              return file.getOriginalFilename();
+            }
+          };
+      bodyBuilder.part("files", resource);
+    }
+    bodyBuilder.part("text", text);
+
+    HttpEntity<MultiValueMap<String, HttpEntity<?>>> requestEntity =
+        new HttpEntity<>(bodyBuilder.build(), headers);
+
+    // Make the POST request to the TTP Extraction AI Webservice
+    ResponseEntity<String> response = restTemplate.postForEntity(url, requestEntity, String.class);
+
+    if (response.getStatusCode().isError()) {
+      log.error("Request to TTP Extraction AI Webservice failed: {}", response.getBody());
+      throw new RestClientException(
+          "Request to TTP Extraction AI Webservice failed: " + response.getBody());
+    }
+    return response.getBody();
+  }
+
+  /**
+   * Extract external attack pattern IDs from the response body of the TTP Extraction AI Webservice.
+   *
+   * @param responseBody The response body from the TTP Extraction AI Webservice, expected to be a
+   *     JSON array
+   * @return Set of external attack pattern IDs extracted from the response
+   * @throws IOException
+   */
+  private Set<String> extractExternalAttackPatternIdsFromResponse(String responseBody)
+      throws IOException {
+    JsonNode fileOrTextJsonArray = mapper.readTree(responseBody);
+    Set<String> externalAttackPatternIds = new HashSet<>();
+
+    // For each (file or text_input) key-value pair in the JSON root
+    for (JsonNode fileOrText : fileOrTextJsonArray) {
+      for (JsonNode chunk : fileOrText) {
+        AnalysisResultFromTTPExtractionAIWebserviceOutput result =
+            mapper.convertValue(chunk, AnalysisResultFromTTPExtractionAIWebserviceOutput.class);
+
+        externalAttackPatternIds.addAll(result.getPredictions().keySet());
+      }
+    }
+    return externalAttackPatternIds;
+  }
+
+  /**
+   * Get the attack pattern IDs from the external IDs.
+   *
+   * @param externalAttackPatternIds Set of external attack pattern IDs to be converted to internal
+   *     IDs.
+   * @return List of attack pattern IDs corresponding to the external IDs.
+   */
+  private List<String> getAttackPatternIds(Set<String> externalAttackPatternIds) {
+    if (!externalAttackPatternIds.isEmpty()) {
+      List<AttackPattern> attackPatterns =
+          this.attackPatternRepository.findAllByExternalIdInIgnoreCase(
+              new ArrayList<>(externalAttackPatternIds));
+      return attackPatterns.stream().map(AttackPattern::getId).toList();
+    }
+    return Collections.emptyList();
+  }
+
+  /**
+   * Validate the inputs for the TTP Extraction AI Webservice.
+   *
+   * @param files List of files to be analyzed, maximum 5 files.
+   * @param text Text input to be analyzed.
+   */
+  private void validateInputs(List<MultipartFile> files, String text) {
+    if (files.isEmpty() && (text == null || text.isBlank())) {
+      throw new IllegalArgumentException("Either files or text must be provided");
+    }
+    if (files.size() > 5) {
+      throw new IllegalArgumentException("Maximum of 5 files allowed");
+    }
+  }
+
+  /**
+   * Search for attack patterns using the TTP Extraction AI Webservice.
+   *
+   * @param files List of files to be analyzed, maximum 5 files.
+   * @param text Text input to be analyzed.
+   * @return List of attack pattern IDs found in the analysis.
+   */
+  public List<String> searchAttackPatternWithTTPAIWebservice(
+      List<MultipartFile> files, String text) {
+    validateInputs(files, text);
+    try {
+      String responseBody = callTTPExtractionAIWebservice(files, text);
+      Set<String> attackPatternExternalIds =
+          extractExternalAttackPatternIdsFromResponse(responseBody);
+      return getAttackPatternIds(attackPatternExternalIds);
+
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+}


Test patch:
diff --git a/openbas-api/src/test/java/io/openbas/rest/attack_pattern/AttackPatternApiTest.java b/openbas-api/src/test/java/io/openbas/rest/attack_pattern/AttackPatternApiTest.java
new file mode 100644
--- /dev/null
+++ b/openbas-api/src/test/java/io/openbas/rest/attack_pattern/AttackPatternApiTest.java
@@ -0,0 +1,148 @@
+package io.openbas.rest.attack_pattern;
+
+import static io.openbas.rest.attack_pattern.AttackPatternApi.ATTACK_PATTERN_URI;
+import static org.junit.jupiter.api.Assertions.*;
+import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS;
+import static org.mockito.ArgumentMatchers.*;
+import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.multipart;
+
+import io.openbas.IntegrationTest;
+import io.openbas.database.model.AttackPattern;
+import io.openbas.database.repository.AttackPatternRepository;
+import io.openbas.ee.Ee;
+import io.openbas.utils.fixtures.files.AttackPatternFixture;
+import io.openbas.utils.mockUser.WithMockAdminUser;
+import jakarta.servlet.ServletException;
+import java.util.Objects;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.mockito.Mockito;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.mock.mockito.MockBean;
+import org.springframework.core.env.Environment;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
+import org.springframework.mock.web.MockMultipartFile;
+import org.springframework.mock.web.MockPart;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.web.servlet.MockMvc;
+import org.springframework.transaction.annotation.Transactional;
+import org.springframework.web.client.RestTemplate;
+
+@Transactional
+@TestInstance(PER_CLASS)
+@DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_CLASS)
+public class AttackPatternApiTest extends IntegrationTest {
+  @Autowired private Environment env;
+
+  @MockBean private RestTemplate mockRestTemplate;
+  @MockBean private Ee mockEe;
+
+  @Autowired private MockMvc mvc;
+  @Autowired private AttackPatternRepository attackPatternRepository;
+
+  @Nested
+  @DisplayName("Search Attack Patterns with AI Webservice")
+  @WithMockAdminUser
+  class AttackPatternWithTTPAIWebservice {
+
+    @Test
+    @DisplayName("Should throw an exception if no text and no files are provided")
+    void given_noTextAndNoFiles_should_throwAnException() {
+      MockPart jsonPart = new MockPart("text", "".getBytes());
+
+      Exception exception =
+          assertThrows(
+              ServletException.class,
+              () ->
+                  mvc.perform(
+                      multipart(ATTACK_PATTERN_URI + "/search-with-ai")
+                          .part(jsonPart)
+                          .contentType(MediaType.MULTIPART_FORM_DATA)));
+
+      assertTrue(exception.getMessage().contains("Either files or text must be provided"));
+    }
+
+    @Test
+    @DisplayName("Should throw an exception if more than 5 files are provided")
+    void given_moreThan5Files_should_throwAnException() {
+      MockPart jsonPart = new MockPart("text", "".getBytes());
+      byte[] content = new byte[] {1, 2, 3, 4, 5}; // Example binary content
+      MockMultipartFile mockFile =
+          new MockMultipartFile("files", "mock-file.pdf", "application/pdf", content);
+      MockMultipartFile mockFile2 =
+          new MockMultipartFile("files", "mock-file.pdf", "application/pdf", content);
+      MockMultipartFile mockFile3 =
+          new MockMultipartFile("files", "mock-file.pdf", "application/pdf", content);
+      MockMultipartFile mockFile4 =
+          new MockMultipartFile("files", "mock-file.pdf", "application/pdf", content);
+      MockMultipartFile mockFile5 =
+          new MockMultipartFile("files", "mock-file.pdf", "application/pdf", content);
+      MockMultipartFile mockFile6 =
+          new MockMultipartFile("files", "mock-file.pdf", "application/pdf", content);
+
+      Exception exception =
+          assertThrows(
+              ServletException.class,
+              () ->
+                  mvc.perform(
+                      multipart(ATTACK_PATTERN_URI + "/search-with-ai")
+                          .part(jsonPart)
+                          .file(mockFile)
+                          .file(mockFile2)
+                          .file(mockFile3)
+                          .file(mockFile4)
+                          .file(mockFile5)
+                          .file(mockFile6)
+                          .contentType(MediaType.MULTIPART_FORM_DATA)));
+
+      assertTrue(exception.getMessage().contains("Maximum of 5 files allowed"));
+    }
+
+    @Test
+    @DisplayName("Should call AI webservice with text and files and return attackPatternsIds")
+    void given_textAndFiles_should_returnAttackPatternsIds() throws Exception {
+      AttackPattern attackPattern = AttackPatternFixture.createDefaultAttackPattern();
+      attackPattern.setExternalId("T1057");
+      AttackPattern attackPatternSaved = attackPatternRepository.save(attackPattern);
+
+      String url = Objects.requireNonNull(env.getProperty("ttp.extraction.ai.webservice.url"));
+      Mockito.when(mockRestTemplate.postForEntity(eq(url), any(), any()))
+          .thenReturn(
+              new ResponseEntity<>(
+                  """
+                          {
+                          "mock-file.pdf": [
+                            {
+                              "text": "Another file's text chunk",
+                              "predictions": {
+                                "T1057": 0.92
+                              }
+                            }
+                          ]
+                        }""",
+                  HttpStatus.OK));
+      Mockito.when(mockEe.getEnterpriseEditionLicensePem()).thenReturn("mock-certificate");
+      MockPart jsonPart = new MockPart("text", "Test".getBytes());
+      byte[] content = new byte[] {1, 2, 3, 4, 5}; // Example binary content
+      MockMultipartFile mockFile =
+          new MockMultipartFile("files", "mock-file.pdf", "application/pdf", content);
+
+      String response =
+          mvc.perform(
+                  multipart(ATTACK_PATTERN_URI + "/search-with-ai")
+                      .part(jsonPart)
+                      .file(mockFile)
+                      .contentType(MediaType.MULTIPART_FORM_DATA))
+              .andReturn()
+              .getResponse()
+              .getContentAsString();
+
+      assertNotNull(response);
+      assertTrue(response.contains(attackPatternSaved.getId()));
+    }
+  }
+}
diff --git a/openbas-api/src/test/resources/application.properties b/openbas-api/src/test/resources/application.properties
--- a/openbas-api/src/test/resources/application.properties
+++ b/openbas-api/src/test/resources/application.properties
@@ -58,6 +58,9 @@ minio.access-secret=minioadmin
 telemetry.obas.endpoint=http://localhost
 telemetry.service.name=openbas-telemetry
 
+## TTP extraction AI Webservice
+ttp.extraction.ai.webservice.url=http://localhost/ttp-extraction
+
 #############
 # INJECTORS #
 #############