Preserve comparison and scan identities

Author: jdalton

socketsecurity/core/__init__.py

@@ -27,6 +27,11 @@
 
 from socketsecurity import USER_AGENT, __version__
 from socketsecurity.core.classes import Alert, Diff, FullScan, Issue, Package, Purl
+from socketsecurity.core.diff_result import (
+    DiffScanResult,
+    PackageComparisonResult,
+    create_diff_scan_result,
+)
 from socketsecurity.core.exceptions import APIResourceNotFound
 
 from .resource_utils import check_file_count_against_ulimit
@@ -1313,7 +1318,7 @@ def create_full_scan_with_report_url(
         base_socket = "https://socket.dev/dashboard/org"
         diff.report_url = f"{base_socket}/{self.config.org_slug}/sbom/{new_full_scan.id}"
         diff.diff_url = diff.report_url
-        diff.id = new_full_scan.id
+        diff.set_scan_ids(new_full_scan.id)
         diff.is_full_scan = True
 
         needs_alerts = (
@@ -1957,12 +1962,12 @@ def get_license_text_via_purl(self, packages: dict[str, Package], batch_size: in
 
         return packages
 
-    def get_diff_scan_artifacts(
+    def _get_diff_scan_result(
             self,
             head_full_scan_id: str,
             new_full_scan_id: str,
             external_href: Optional[str] = None
-    ) -> DiffArtifacts:
+    ) -> DiffScanResult:
         """Compare two full scans via the diff-scans endpoints, polling for the result.
 
         Creates a diff-scan resource from the two full scan IDs, then polls
@@ -2110,10 +2115,19 @@ def get_diff_scan_artifacts(
             last_interval = interval
             interval = min(interval * DIFF_SCAN_POLL_BACKOFF_MULTIPLIER, DIFF_SCAN_POLL_MAX_INTERVAL_SECONDS)
 
-        return DiffArtifacts.from_dict({
-            key: artifacts_dict.get(key) or []
-            for key in ("added", "removed", "unchanged", "replaced", "updated")
-        })
+        return create_diff_scan_result(artifacts_dict, diff_scan_id)
+
+    def get_diff_scan_artifacts(
+            self,
+            head_full_scan_id: str,
+            new_full_scan_id: str,
+            external_href: Optional[str] = None
+    ) -> DiffArtifacts:
+        return self._get_diff_scan_result(
+            head_full_scan_id,
+            new_full_scan_id,
+            external_href=external_href,
+        ).artifacts
 
     def _requires_unchanged_artifacts(self) -> bool:
         """Whether any enabled output reads the unchanged half of a comparison.
@@ -2148,13 +2162,13 @@ def _requires_unchanged_artifacts(self) -> bool:
             or getattr(config, "legal_format", "socket") == "fossa"
         )
 
-    def get_added_and_removed_packages(
+    def _get_package_comparison(
             self,
             head_full_scan_id: str,
             new_full_scan_id: str,
             include_license_details: bool = False,
             external_href: Optional[str] = None
-    ) -> Tuple[Dict[str, Package], Dict[str, Package], Dict[str, Package]]:
+    ) -> PackageComparisonResult:
         """
         Get packages that were added and removed between scans.
 
@@ -2190,14 +2204,14 @@ def get_added_and_removed_packages(
                 with the primary diff-scan resource
 
         Returns:
-            Tuple of (added_packages, removed_packages) dictionaries
+            Added, removed, and all-package maps plus the completed diff ID
         """
 
         log.info(f"Comparing scans - Head scan ID: {head_full_scan_id}, New scan ID: {new_full_scan_id}")
         diff_start = time.time()
-        diff_artifacts = None
+        comparison = None
         try:
-            diff_artifacts = self.get_diff_scan_artifacts(
+            comparison = self._get_diff_scan_result(
                 head_full_scan_id,
                 new_full_scan_id,
                 external_href=external_href,
@@ -2215,7 +2229,7 @@ def get_added_and_removed_packages(
                 f"reason={type(error).__name__}"
             )
 
-        if diff_artifacts is None:
+        if comparison is None:
             try:
                 diff_artifacts = (
                     self.sdk.fullscans.stream_diff(
@@ -2226,6 +2240,10 @@ def get_added_and_removed_packages(
                         include_license_details=str(include_license_details).lower()
                     ).data.artifacts
                 )
+                comparison = DiffScanResult(
+                    artifacts=diff_artifacts,
+                    diff_scan_id=None,
+                )
             except APIFailure as e:
                 log.error(f"API Error: {e}")
                 # API failures are mapped to the configured infrastructure exit code by cli().
@@ -2245,7 +2263,7 @@ def get_added_and_removed_packages(
         marker_found = False
         buckets: Dict[str, List] = {}
         for name in ("added", "removed", "unchanged", "replaced", "updated"):
-            bucket = getattr(diff_artifacts, name)
+            bucket = getattr(comparison.artifacts, name)
             buckets[name] = [a for a in bucket if not Core.is_invalid_facts_marker(a)]
             marker_found = marker_found or len(buckets[name]) != len(bucket)
         Core.warn_if_invalid_facts_marker(marker_found)
@@ -2304,7 +2322,27 @@ def get_added_and_removed_packages(
         else:
             log.debug("Skipping PURL endpoint call (--generate-license not set)")
 
-        return added_packages, removed_packages, packages
+        return PackageComparisonResult(
+            added_packages=added_packages,
+            removed_packages=removed_packages,
+            packages=packages,
+            diff_scan_id=comparison.diff_scan_id,
+        )
+
+    def get_added_and_removed_packages(
+            self,
+            head_full_scan_id: str,
+            new_full_scan_id: str,
+            include_license_details: bool = False,
+            external_href: Optional[str] = None
+    ) -> Tuple[Dict[str, Package], Dict[str, Package], Dict[str, Package]]:
+        result = self._get_package_comparison(
+            head_full_scan_id,
+            new_full_scan_id,
+            include_license_details=include_license_details,
+            external_href=external_href,
+        )
+        return result.added_packages, result.removed_packages, result.packages
 
     def create_new_diff(
             self,
@@ -2456,16 +2494,15 @@ def create_new_diff(
         # the response and risks the truncation crash on large repos. The
         # user flag still controls the dashboard report URL below; it just no
         # longer gates this internal diff payload.
-        (
-            added_packages,
-            removed_packages,
-            packages
-        ) = self.get_added_and_removed_packages(
+        comparison = self._get_package_comparison(
             head_full_scan_id,
             new_full_scan.id,
             include_license_details=False,
             external_href=external_href,
         )
+        added_packages = comparison.added_packages
+        removed_packages = comparison.removed_packages
+        packages = comparison.packages
 
         # Separate unchanged packages from added/removed for --strict-blocking support
         unchanged_packages = {
@@ -2477,13 +2514,16 @@ def create_new_diff(
         diff.packages = packages
 
         base_socket = "https://socket.dev/dashboard/org"
-        diff.id = new_full_scan.id
+        diff.set_scan_ids(
+            new_full_scan.id,
+            old_scan_id=head_full_scan_id,
+            diff_scan_id=comparison.diff_scan_id,
+        )
 
         report_url = f"{base_socket}/{self.config.org_slug}/sbom/{diff.id}"
         if not params.include_license_details:
             report_url += "?include_license_details=false"
         diff.report_url = report_url
-        diff.new_scan_id = new_full_scan.id
 
         if head_full_scan_id is not None:
             diff.diff_url = f"{base_socket}/{self.config.org_slug}/diff/{head_full_scan_id}/{diff.id}"

socketsecurity/core/alert_selection.py

@@ -38,7 +38,9 @@ def clone_diff_with_selected_alerts(diff: Diff, selected_alerts: List[Issue]) ->
     )
     selected_diff.id = getattr(diff, "id", "")
     selected_diff.report_url = getattr(diff, "report_url", "")
-    selected_diff.new_scan_id = getattr(diff, "new_scan_id", "")
+    selected_diff.new_scan_id = getattr(diff, "new_scan_id", None)
+    selected_diff.old_scan_id = getattr(diff, "old_scan_id", None)
+    selected_diff.diff_scan_id = getattr(diff, "diff_scan_id", None)
     return selected_diff
 
 

socketsecurity/core/classes.py

@@ -571,7 +571,9 @@ class Diff:
     sbom: str
     report_url: str
     diff_url: str
-    new_scan_id: str
+    new_scan_id: Optional[str]
+    old_scan_id: Optional[str]
+    diff_scan_id: Optional[str]
     is_full_scan: bool
     alerts_fetched: bool
 
@@ -599,6 +601,24 @@ def __init__(self, **kwargs):
             self.is_full_scan = False
         if not hasattr(self, "alerts_fetched"):
             self.alerts_fetched = False
+        if not hasattr(self, "new_scan_id"):
+            self.new_scan_id = None
+        if not hasattr(self, "old_scan_id"):
+            self.old_scan_id = None
+        if not hasattr(self, "diff_scan_id"):
+            self.diff_scan_id = None
+
+    def set_scan_ids(
+        self,
+        new_scan_id: str,
+        *,
+        old_scan_id: Optional[str] = None,
+        diff_scan_id: Optional[str] = None,
+    ) -> None:
+        self.id = new_scan_id
+        self.new_scan_id = new_scan_id
+        self.old_scan_id = old_scan_id
+        self.diff_scan_id = diff_scan_id
 
     def __str__(self):
         return json.dumps(self.__dict__)

socketsecurity/core/diff_result.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from socketdev.fullscans import DiffArtifacts
+
+from socketsecurity.core.classes import Package
+
+
+@dataclass
+class DiffScanResult:
+    artifacts: DiffArtifacts
+    diff_scan_id: Optional[str]
+
+
+@dataclass
+class PackageComparisonResult:
+    added_packages: Dict[str, Package]
+    removed_packages: Dict[str, Package]
+    packages: Dict[str, Package]
+    diff_scan_id: Optional[str]
+
+
+def create_diff_scan_result(
+    artifacts: Dict[str, object], diff_scan_id: Optional[str]
+) -> DiffScanResult:
+    return DiffScanResult(
+        artifacts=DiffArtifacts.from_dict({
+            key: artifacts.get(key) or []
+            for key in ("added", "removed", "unchanged", "replaced", "updated")
+        }),
+        diff_scan_id=diff_scan_id,
+    )