Merge pull request #63 from tls-attacker/feature/add-scanjobdescription-to-scan-result

Add scanJobDescription ID and timestamp to ScanResult with retrieval method

Author: ic0ns

src/main/java/de/rub/nds/crawler/data/ScanConfig.java

@@ -11,7 +11,6 @@
 import de.rub.nds.crawler.core.BulkScanWorker;
 import de.rub.nds.scanner.core.config.ScannerDetail;
 import de.rub.nds.scanner.core.probe.ProbeType;
-
 import java.io.Serializable;
 import java.util.List;
 
@@ -30,15 +29,14 @@ public abstract class ScanConfig implements Serializable {
     private List<ProbeType> excludedProbes;
 
     @SuppressWarnings("unused")
-    private ScanConfig() {
-    }
+    private ScanConfig() {}
 
     /**
      * Creates a new scan configuration with the specified parameters.
      *
      * @param scannerDetail The level of detail for the scan
-     * @param reexecutions  The number of times to retry failed scans
-     * @param timeout       The timeout for each scan in seconds
+     * @param reexecutions The number of times to retry failed scans
+     * @param timeout The timeout for each scan in seconds
      */
     protected ScanConfig(ScannerDetail scannerDetail, int reexecutions, int timeout) {
         this(scannerDetail, reexecutions, timeout, null);
@@ -121,9 +119,9 @@ public void setExcludedProbes(List<ProbeType> excludedProbes) {
      * Creates a worker for this scan configuration. Each implementation must provide a factory
      * method to create the appropriate worker type.
      *
-     * @param bulkScanID                The ID of the bulk scan this worker is for
+     * @param bulkScanID The ID of the bulk scan this worker is for
      * @param parallelConnectionThreads The number of parallel connection threads to use
-     * @param parallelScanThreads       The number of parallel scan threads to use
+     * @param parallelScanThreads The number of parallel scan threads to use
      * @return A worker for this scan configuration
      */
     public abstract BulkScanWorker<? extends ScanConfig> createWorker(

src/main/java/de/rub/nds/crawler/data/ScanResult.java

@@ -12,6 +12,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import de.rub.nds.crawler.constant.JobStatus;
 import java.io.Serializable;
+import java.time.Instant;
 import java.util.UUID;
 import org.bson.Document;
 
@@ -27,25 +28,35 @@ public class ScanResult implements Serializable {
 
     private final Document result;
 
+    private final String scanJobDescriptionId;
+
+    private final Instant timestamp;
+
     @JsonCreator
     private ScanResult(
+            @JsonProperty("scanJobDescription") String scanJobDescriptionId,
             @JsonProperty("bulkScan") String bulkScan,
             @JsonProperty("scanTarget") ScanTarget scanTarget,
             @JsonProperty("resultStatus") JobStatus jobStatus,
-            @JsonProperty("result") Document result) {
+            @JsonProperty("result") Document result,
+            @JsonProperty("timestamp") Instant timestamp) {
         this.id = UUID.randomUUID().toString();
+        this.scanJobDescriptionId = scanJobDescriptionId;
         this.bulkScan = bulkScan;
         this.scanTarget = scanTarget;
         this.jobStatus = jobStatus;
         this.result = result;
+        this.timestamp = timestamp != null ? timestamp : Instant.now();
     }
 
     public ScanResult(ScanJobDescription scanJobDescription, Document result) {
         this(
+                scanJobDescription.getId().toString(),
                 scanJobDescription.getBulkScanInfo().getBulkScanId(),
                 scanJobDescription.getScanTarget(),
                 scanJobDescription.getStatus(),
-                result);
+                result,
+                Instant.now());
         if (scanJobDescription.getStatus() == JobStatus.TO_BE_EXECUTED) {
             throw new IllegalArgumentException(
                     "ScanJobDescription must not be in TO_BE_EXECUTED state");
@@ -86,4 +97,14 @@ public Document getResult() {
     public JobStatus getResultStatus() {
         return jobStatus;
     }
+
+    @JsonProperty("scanJobDescription")
+    public String getScanJobDescriptionId() {
+        return scanJobDescriptionId;
+    }
+
+    @JsonProperty("timestamp")
+    public Instant getTimestamp() {
+        return timestamp;
+    }
 }

src/main/java/de/rub/nds/crawler/persistence/IPersistenceProvider.java

@@ -61,4 +61,16 @@ public interface IPersistenceProvider {
      * @return The scan result, or null if not found.
      */
     ScanResult getScanResultById(String dbName, String collectionName, String id);
+
+    /**
+     * Retrieve the most recent scan result by its scan job description ID. If multiple results
+     * exist for the same scan job description ID, returns the one with the latest timestamp.
+     *
+     * @param dbName The database name where the scan result is stored.
+     * @param collectionName The collection name where the scan result is stored.
+     * @param scanJobDescriptionId The scan job description ID to search for.
+     * @return The most recent scan result, or null if not found.
+     */
+    ScanResult getScanResultByScanJobDescriptionId(
+            String dbName, String collectionName, String scanJobDescriptionId);
 }

src/main/java/de/rub/nds/crawler/persistence/MongoPersistenceProvider.java

@@ -205,6 +205,8 @@ private JacksonMongoCollection<ScanResult> initResultCollection(
         collection.createIndex(Indexes.ascending("scanTarget.hostname"));
         collection.createIndex(Indexes.ascending("scanTarget.trancoRank"));
         collection.createIndex(Indexes.ascending("scanTarget.resultStatus"));
+        collection.createIndex(Indexes.ascending("scanJobDescription"));
+        collection.createIndex(Indexes.descending("timestamp"));
         return collection;
     }
 
@@ -343,4 +345,51 @@ public ScanResult getScanResultById(String dbName, String collectionName, String
             throw new RuntimeException("Failed to retrieve scan result with ID: " + id, e);
         }
     }
+
+    @Override
+    public ScanResult getScanResultByScanJobDescriptionId(
+            String dbName, String collectionName, String scanJobDescriptionId) {
+        LOGGER.info(
+                "Retrieving most recent scan result for scanJobDescriptionId {} from collection: {}.{}",
+                scanJobDescriptionId,
+                dbName,
+                collectionName);
+
+        try {
+            var collection = resultCollectionCache.getUnchecked(Pair.of(dbName, collectionName));
+
+            var query = new org.bson.Document("scanJobDescription", scanJobDescriptionId);
+
+            var iterable = collection.find(query).sort(new org.bson.Document("timestamp", -1));
+
+            var iterator = iterable.iterator();
+            ScanResult result = null;
+            if (iterator.hasNext()) {
+                result = iterator.next();
+            }
+
+            if (result == null) {
+                LOGGER.warn(
+                        "No scan result found for scanJobDescriptionId: {} in collection: {}.{}",
+                        scanJobDescriptionId,
+                        dbName,
+                        collectionName);
+            } else {
+                LOGGER.info(
+                        "Retrieved most recent scan result for scanJobDescriptionId: {} from collection: {}.{} (timestamp: {})",
+                        scanJobDescriptionId,
+                        dbName,
+                        collectionName,
+                        result.getTimestamp());
+            }
+
+            return result;
+        } catch (Exception e) {
+            LOGGER.error("Exception while retrieving scan result from MongoDB: ", e);
+            throw new RuntimeException(
+                    "Failed to retrieve scan result for scanJobDescriptionId: "
+                            + scanJobDescriptionId,
+                    e);
+        }
+    }
 }

src/test/java/de/rub/nds/crawler/dummy/DummyPersistenceProvider.java

@@ -46,4 +46,13 @@ public ScanResult getScanResultById(String dbName, String collectionName, String
                 .findFirst()
                 .orElse(null);
     }
+
+    @Override
+    public ScanResult getScanResultByScanJobDescriptionId(
+            String dbName, String collectionName, String id) {
+        return results.stream()
+                .filter(result -> result.getScanJobDescriptionId().equals(id))
+                .max((r1, r2) -> r1.getTimestamp().compareTo(r2.getTimestamp()))
+                .orElse(null);
+    }
 }

src/test/java/de/rub/nds/crawler/dummy/DummyPersistenceProviderTest.java

@@ -0,0 +1,101 @@
+/*
+ * TLS-Crawler - A TLS scanning tool to perform large scale scans with the TLS-Scanner
+ *
+ * Copyright 2018-2023 Ruhr University Bochum, Paderborn University, and Hackmanit GmbH
+ *
+ * Licensed under Apache License, Version 2.0
+ * http://www.apache.org/licenses/LICENSE-2.0.txt
+ */
+package de.rub.nds.crawler.dummy;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import de.rub.nds.crawler.constant.JobStatus;
+import de.rub.nds.crawler.core.BulkScanWorker;
+import de.rub.nds.crawler.data.BulkScan;
+import de.rub.nds.crawler.data.ScanConfig;
+import de.rub.nds.crawler.data.ScanJobDescription;
+import de.rub.nds.crawler.data.ScanResult;
+import de.rub.nds.crawler.data.ScanTarget;
+import de.rub.nds.scanner.core.config.ScannerDetail;
+import org.bson.Document;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class DummyPersistenceProviderTest {
+
+    private DummyPersistenceProvider provider;
+    private BulkScan testBulkScan;
+
+    @BeforeEach
+    void setUp() {
+        provider = new DummyPersistenceProvider();
+
+        // Create a test BulkScan
+        ScanConfig scanConfig =
+                new ScanConfig(ScannerDetail.NORMAL, 1, 5000) {
+                    @Override
+                    public BulkScanWorker<? extends ScanConfig> createWorker(
+                            String bulkScanID,
+                            int parallelConnectionThreads,
+                            int parallelScanThreads) {
+                        return null;
+                    }
+                };
+
+        testBulkScan =
+                new BulkScan(
+                        this.getClass(),
+                        this.getClass(),
+                        "test-scan",
+                        scanConfig,
+                        System.currentTimeMillis(),
+                        false,
+                        null);
+        testBulkScan.set_id("test-bulk-scan-id");
+    }
+
+    @Test
+    void testGetScanResultByScanJobDescriptionId_ReturnsMostRecent() throws InterruptedException {
+        ScanTarget target = new ScanTarget();
+        target.setHostname("example.com");
+        target.setIp("93.184.216.34");
+        target.setPort(443);
+
+        ScanJobDescription jobDescription =
+                new ScanJobDescription(target, testBulkScan, JobStatus.SUCCESS);
+        String scanJobDescriptionId = jobDescription.getId().toString();
+
+        Document resultDoc1 = new Document();
+        resultDoc1.put("attempt", 1);
+        ScanResult scanResult1 = new ScanResult(jobDescription, resultDoc1);
+        provider.insertScanResult(scanResult1, jobDescription);
+
+        Thread.sleep(10);
+
+        Document resultDoc2 = new Document();
+        resultDoc2.put("attempt", 2);
+        ScanResult scanResult2 = new ScanResult(jobDescription, resultDoc2);
+        provider.insertScanResult(scanResult2, jobDescription);
+
+        Thread.sleep(10);
+
+        Document resultDoc3 = new Document();
+        resultDoc3.put("attempt", 3);
+        ScanResult scanResult3 = new ScanResult(jobDescription, resultDoc3);
+        provider.insertScanResult(scanResult3, jobDescription);
+
+        ScanResult retrieved =
+                provider.getScanResultByScanJobDescriptionId(
+                        "test-db", "test-collection", scanJobDescriptionId);
+
+        assertNotNull(retrieved);
+        assertEquals(scanJobDescriptionId, retrieved.getScanJobDescriptionId());
+
+        assertTrue(retrieved.getTimestamp().compareTo(scanResult1.getTimestamp()) >= 0);
+        assertTrue(retrieved.getTimestamp().compareTo(scanResult2.getTimestamp()) >= 0);
+        assertTrue(retrieved.getTimestamp().compareTo(scanResult3.getTimestamp()) >= 0);
+
+        assertEquals(scanResult3.getTimestamp(), retrieved.getTimestamp());
+    }
+}