From 2013216d25c78aecbc3c0273a87050c757d1347b Mon Sep 17 00:00:00 2001
From: Ivan Milchev <ivan@mondoo.com>
Date: Sat, 17 Feb 2024 01:25:23 +0100
Subject: [PATCH] fix shell command

Signed-off-by: Ivan Milchev <ivan@mondoo.com>
---
 apps/cnquery/cmd/shell.go     | 23 +++++++++++------------
 cli/components/assetselect.go |  6 +++---
 explorer/scan/discovery.go    |  6 +++---
 3 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/apps/cnquery/cmd/shell.go b/apps/cnquery/cmd/shell.go
index d8faaede13..b2b228ddc1 100644
--- a/apps/cnquery/cmd/shell.go
+++ b/apps/cnquery/cmd/shell.go
@@ -108,11 +108,17 @@ func StartShell(runtime *providers.Runtime, conf *ShellConfig) error {
 
 	connectAsset := filteredAssets[0]
 	if len(filteredAssets) > 1 {
+		invAssets := make([]*inventory.Asset, 0, len(filteredAssets))
+		for _, a := range filteredAssets {
+			invAssets = append(invAssets, a.Asset)
+		}
+
 		isTTY := isatty.IsTerminal(os.Stdout.Fd())
 		if isTTY {
-			connectAsset = components.AssetSelect(filteredAssets)
+			selectedAsset := components.AssetSelect(invAssets)
+			connectAsset = filteredAssets[selectedAsset]
 		} else {
-			fmt.Println(components.AssetList(theme.OperatingSystemTheme, filteredAssets))
+			fmt.Println(components.AssetList(theme.OperatingSystemTheme, invAssets))
 			log.Fatal().Msg("cannot connect to more than one asset, use --platform-id to select a specific asset")
 		}
 	}
@@ -121,19 +127,12 @@ func StartShell(runtime *providers.Runtime, conf *ShellConfig) error {
 		log.Fatal().Msg("no asset selected")
 	}
 
-	err = runtime.Connect(&plugin.ConnectReq{
-		Features: conf.Features,
-		Asset:    connectAsset,
-		Upstream: conf.UpstreamConfig,
-	})
-	if err != nil {
-		log.Fatal().Err(err).Msg("failed to connect to asset")
-	}
-	log.Info().Msgf("connected to %s", runtime.Provider.Connection.Asset.Platform.Title)
+	log.Info().Msgf("connected to %s", connectAsset.Runtime.Provider.Connection.Asset.Platform.Title)
 
 	// when we close the shell, we need to close the backend and store the recording
 	onCloseHandler := func() {
 		runtime.Close()
+		connectAsset.Runtime.Close()
 		providers.Coordinator.Shutdown()
 	}
 
@@ -142,7 +141,7 @@ func StartShell(runtime *providers.Runtime, conf *ShellConfig) error {
 	shellOptions = append(shellOptions, shell.WithFeatures(conf.Features))
 	shellOptions = append(shellOptions, shell.WithUpstreamConfig(conf.UpstreamConfig))
 
-	sh, err := shell.New(runtime, shellOptions...)
+	sh, err := shell.New(connectAsset.Runtime, shellOptions...)
 	if err != nil {
 		log.Error().Err(err).Msg("failed to initialize interactive shell")
 	}
diff --git a/cli/components/assetselect.go b/cli/components/assetselect.go
index fd5180e6aa..57d36cbfec 100644
--- a/cli/components/assetselect.go
+++ b/cli/components/assetselect.go
@@ -12,7 +12,7 @@ import (
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/inventory"
 )
 
-func AssetSelect(assetList []*inventory.Asset) *inventory.Asset {
+func AssetSelect(assetList []*inventory.Asset) int {
 	list := make([]string, len(assetList))
 
 	// map asset name to list
@@ -36,9 +36,9 @@ func AssetSelect(assetList []*inventory.Asset) *inventory.Asset {
 	}
 
 	if selection == -1 {
-		return nil
+		return -1
 	}
 	selected := assetList[selection]
 	log.Info().Int("selection", selection).Str("asset", selected.Name).Msg("selected asset")
-	return selected
+	return selection
 }
diff --git a/explorer/scan/discovery.go b/explorer/scan/discovery.go
index 9794f17feb..a3711aaa08 100644
--- a/explorer/scan/discovery.go
+++ b/explorer/scan/discovery.go
@@ -57,12 +57,12 @@ func (d *DiscoveredAssets) AddError(asset *inventory.Asset, err error) {
 	d.Errors = append(d.Errors, &AssetWithError{Asset: asset, Err: err})
 }
 
-func (d *DiscoveredAssets) GetAssetsByPlatformID(platformID string) []*inventory.Asset {
-	var assets []*inventory.Asset
+func (d *DiscoveredAssets) GetAssetsByPlatformID(platformID string) []*AssetWithRuntime {
+	var assets []*AssetWithRuntime
 	for _, a := range d.Assets {
 		for _, p := range a.Asset.PlatformIds {
 			if platformID == "" || p == platformID {
-				assets = append(assets, a.Asset)
+				assets = append(assets, a)
 				break
 			}
 		}