From 3958317420507f403c11f2b320cc0521c609b379 Mon Sep 17 00:00:00 2001
From: Ivan Milchev <ivan@mondoo.com>
Date: Wed, 3 Jan 2024 16:16:55 +0200
Subject: [PATCH] retrieve space bundle instead of asset bundle for upstream
 scans

Signed-off-by: Ivan Milchev <ivan@mondoo.com>
---
 policy/scan/local_scanner.go | 39 +++++++++++++++++++++++++++++++++---
 policy/scan/scan.go          |  1 +
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/policy/scan/local_scanner.go b/policy/scan/local_scanner.go
index a101d05f..29ca8c9b 100644
--- a/policy/scan/local_scanner.go
+++ b/policy/scan/local_scanner.go
@@ -416,6 +416,23 @@ func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *up
 		}
 	}()
 
+	// Retrieve the space bundle if we are not running in incognito mode
+	var spaceBundleMap *policy.PolicyBundleMap
+	if upstream != nil && upstream.ApiEndpoint != "" && !upstream.Incognito {
+		client, err := upstream.InitClient()
+		if err != nil {
+			return nil, err
+		}
+		services, err := policy.NewRemoteServices(client.ApiEndpoint, client.Plugins, client.HttpClient)
+		if err != nil {
+			return nil, err
+		}
+		spaceBundle, err := services.GetBundle(ctx, &policy.Mrn{Mrn: client.SpaceMrn})
+		if err == nil {
+			spaceBundleMap = spaceBundle.ToMap()
+		}
+	}
+
 	assetBatches := batch(assets, 100)
 	for i := range assetBatches {
 		batch := assetBatches[i]
@@ -519,6 +536,7 @@ func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *up
 					UpstreamConfig:   upstream,
 					Asset:            asset,
 					Bundle:           job.Bundle,
+					SpaceBundleMap:   spaceBundleMap,
 					Props:            job.Props,
 					PolicyFilters:    preprocessPolicyFilters(job.PolicyFilters),
 					Ctx:              ctx,
@@ -637,6 +655,12 @@ func (s *LocalScanner) runMotorizedAsset(job *AssetJob) (*AssetReport, error) {
 			if err != nil {
 				return err
 			}
+			// If we have a space bundle at this point, make sure it is in the local cache for the asset
+			if job.SpaceBundleMap != nil {
+				if err := services.SetBundleMap(job.Ctx, job.SpaceBundleMap); err != nil {
+					return err
+				}
+			}
 			services.Upstream = upstream
 		}
 
@@ -984,15 +1008,24 @@ func (s *localAssetScanner) runPolicy() (*policy.Bundle, *policy.ResolvedPolicy,
 	var resolver policy.PolicyResolver = s.services
 
 	log.Debug().Str("asset", s.job.Asset.Mrn).Msg("client> request policies bundle for asset")
-	assetBundle, err := hub.GetBundle(s.job.Ctx, &policy.Mrn{Mrn: s.job.Asset.Mrn})
+
+	// For non-incognito scans we use the space bundle since it contains all queries and controls.
+	// Only exceptions are defined on per-asset basis. Exceptions aren't relevant in this context,
+	// so there is no need to retrieve the asset bundle.
+	bundleMrn := s.job.Asset.Mrn
+	if !s.job.UpstreamConfig.Incognito {
+		bundleMrn = s.job.UpstreamConfig.SpaceMrn
+	}
+
+	assetBundle, err := hub.GetBundle(s.job.Ctx, &policy.Mrn{Mrn: bundleMrn})
 	if err != nil {
 		return nil, nil, err
 	}
 	log.Debug().Msg("client> got policy bundle")
 	logger.TraceJSON(assetBundle)
-	logger.DebugDumpJSON("assetBundle", assetBundle)
+	logger.DebugDumpYAML("spaceBundle", assetBundle)
 
-	rawFilters, err := hub.GetPolicyFilters(s.job.Ctx, &policy.Mrn{Mrn: s.job.Asset.Mrn})
+	rawFilters, err := hub.GetPolicyFilters(s.job.Ctx, &policy.Mrn{Mrn: bundleMrn})
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/policy/scan/scan.go b/policy/scan/scan.go
index 08d2f399..c447f27c 100644
--- a/policy/scan/scan.go
+++ b/policy/scan/scan.go
@@ -32,6 +32,7 @@ type AssetJob struct {
 	UpstreamConfig   *upstream.UpstreamConfig
 	Asset            *inventory.Asset
 	Bundle           *policy.Bundle
+	SpaceBundleMap   *policy.PolicyBundleMap
 	PolicyFilters    []string
 	Props            map[string]string
 	Ctx              context.Context