Merge pull request #2168 from BishopFox/feature/mutli-console

Feature/mutli console

Author: moloch--

client/assets/settings.go

@@ -22,6 +22,7 @@ import (
 	"encoding/json"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"gopkg.in/yaml.v3"
 )
@@ -31,6 +32,18 @@ const (
 	settingsLegacyFileName = "tui-settings.json"
 )
 
+const (
+	// Accepted values for tui-settings.yaml "prompt:":
+	// - host: show "[host]" prefix on sliver-client, "[server]" on server console
+	// - operator-host: show "[operator@host]" prefix on sliver-client, "[server]" on server console
+	// - basic: show "sliver >" (no prefix)
+	// - custom: render full prompt from prompt_template
+	PromptStyleHost         = "host"
+	PromptStyleOperatorHost = "operator-host"
+	PromptStyleBasic        = "basic"
+	PromptStyleCustom       = "custom"
+)
+
 // ClientSettings - Client JSON config
 type ClientSettings struct {
 	TableStyle        string `json:"tables" yaml:"tables"`
@@ -41,6 +54,8 @@ type ClientSettings struct {
 	VimMode           bool   `json:"vim_mode" yaml:"vim_mode"`
 	UserConnect       bool   `json:"user_connect" yaml:"user_connect"`
 	ConsoleLogs       bool   `json:"console_logs" yaml:"console_logs"`
+	PromptStyle       string `json:"prompt" yaml:"prompt"`
+	PromptTemplate    string `json:"prompt_template" yaml:"prompt_template"`
 }
 
 // LoadSettings - Load the client settings from disk
@@ -67,6 +82,8 @@ func LoadSettings() (*ClientSettings, error) {
 		return defaultSettings(), err
 	}
 
+	// Ensure any missing/unknown values are coerced to a supported prompt style.
+	settings.PromptStyle = NormalizePromptStyle(settings.PromptStyle)
 	if err := SaveSettings(settings); err != nil {
 		return settings, err
 	}
@@ -78,6 +95,38 @@ func LoadSettings() (*ClientSettings, error) {
 	return settings, nil
 }
 
+// NormalizePromptStyle canonicalizes prompt style strings and returns a safe
+// default if the value is empty/unknown.
+func NormalizePromptStyle(v string) string {
+	switch strings.ToLower(strings.TrimSpace(v)) {
+	case PromptStyleOperatorHost:
+		return PromptStyleOperatorHost
+	case PromptStyleHost:
+		return PromptStyleHost
+	case PromptStyleBasic:
+		return PromptStyleBasic
+	case PromptStyleCustom:
+		return PromptStyleCustom
+
+	// Backward compatible aliases.
+	case "show host":
+		return PromptStyleHost
+	case "show user and host":
+		return PromptStyleOperatorHost
+	case "show operator and host":
+		return PromptStyleOperatorHost
+	case "operator@host":
+		return PromptStyleOperatorHost
+	case "user@host":
+		return PromptStyleOperatorHost
+
+	case "":
+		return PromptStyleHost
+	default:
+		return PromptStyleHost
+	}
+}
+
 func defaultSettings() *ClientSettings {
 	return &ClientSettings{
 		TableStyle:        "SliverDefault",
@@ -87,9 +136,13 @@ func defaultSettings() *ClientSettings {
 		AlwaysOverflow:    false,
 		VimMode:           false,
 		ConsoleLogs:       true,
+		PromptStyle:       PromptStyleHost,
+		PromptTemplate:    DefaultPromptTemplate,
 	}
 }
 
+const DefaultPromptTemplate = `{{- if .IsServer -}}{{ .Styles.Bold.Render "[server]" }} {{ .Styles.Underline.Render "sliver" }}{{ .Target.Suffix }} > {{- else -}}{{- if .Host -}}{{ .Styles.BoldPrimary.Render (printf "[%s]" .Host) }} {{- end -}}{{ .Styles.Underline.Render "sliver" }}{{ .Target.Suffix }} > {{- end -}}`
+
 // SaveSettings - Save the current settings to disk
 func SaveSettings(settings *ClientSettings) error {
 	rootDir, _ := filepath.Abs(GetRootAppDir())

client/assets/settings_test.go

@@ -21,6 +21,8 @@ always_overflow: true
 vim_mode: true
 user_connect: true
 console_logs: false
+prompt: "basic"
+prompt_template: "{{.Host}}"
 `)
 	if err := os.WriteFile(settingsPath, data, 0o600); err != nil {
 		t.Fatalf("failed to write yaml settings: %v", err)
@@ -54,6 +56,12 @@ console_logs: false
 	if settings.ConsoleLogs {
 		t.Fatalf("expected ConsoleLogs false")
 	}
+	if settings.PromptStyle != PromptStyleBasic {
+		t.Fatalf("expected PromptStyle %q, got %q", PromptStyleBasic, settings.PromptStyle)
+	}
+	if settings.PromptTemplate != "{{.Host}}" {
+		t.Fatalf("expected PromptTemplate %q, got %q", "{{.Host}}", settings.PromptTemplate)
+	}
 }
 
 func TestLoadSettingsMigratesLegacyJSON(t *testing.T) {
@@ -71,6 +79,8 @@ func TestLoadSettingsMigratesLegacyJSON(t *testing.T) {
 		VimMode:           true,
 		UserConnect:       true,
 		ConsoleLogs:       false,
+		PromptStyle:       PromptStyleOperatorHost,
+		PromptTemplate:    "{{.Operator}}@{{.Host}} sliver > ",
 	}
 	data, err := json.MarshalIndent(legacy, "", "  ")
 	if err != nil {
@@ -87,6 +97,12 @@ func TestLoadSettingsMigratesLegacyJSON(t *testing.T) {
 	if settings.TableStyle != "Legacy" {
 		t.Fatalf("expected TableStyle %q, got %q", "Legacy", settings.TableStyle)
 	}
+	if settings.PromptStyle != PromptStyleOperatorHost {
+		t.Fatalf("expected PromptStyle %q, got %q", PromptStyleOperatorHost, settings.PromptStyle)
+	}
+	if settings.PromptTemplate != "{{.Operator}}@{{.Host}} sliver > " {
+		t.Fatalf("expected PromptTemplate %q, got %q", "{{.Operator}}@{{.Host}} sliver > ", settings.PromptTemplate)
+	}
 	if _, err := os.Stat(filepath.Join(rootDir, settingsFileName)); err != nil {
 		t.Fatalf("expected yaml settings to exist: %v", err)
 	}
@@ -108,6 +124,12 @@ func TestLoadSettingsWritesDefault(t *testing.T) {
 	if settings.TableStyle != "SliverDefault" {
 		t.Fatalf("expected default TableStyle %q, got %q", "SliverDefault", settings.TableStyle)
 	}
+	if settings.PromptStyle != PromptStyleHost {
+		t.Fatalf("expected default PromptStyle %q, got %q", PromptStyleHost, settings.PromptStyle)
+	}
+	if settings.PromptTemplate != DefaultPromptTemplate {
+		t.Fatalf("expected default PromptTemplate %q, got %q", DefaultPromptTemplate, settings.PromptTemplate)
+	}
 	rootDir, _ := filepath.Abs(GetRootAppDir())
 	if _, err := os.Stat(filepath.Join(rootDir, settingsFileName)); err != nil {
 		t.Fatalf("expected default settings to be written: %v", err)

client/cli/config.go

@@ -26,16 +26,16 @@ import (
 	"github.com/bishopfox/sliver/client/forms"
 )
 
-func selectConfig() *assets.ClientConfig {
+func selectConfig() (string, *assets.ClientConfig) {
 	configs := assets.GetConfigs()
 
 	if len(configs) == 0 {
-		return nil
+		return "", nil
 	}
 
 	if len(configs) == 1 {
-		for _, config := range configs {
-			return config
+		for key, config := range configs {
+			return key, config
 		}
 	}
 
@@ -49,8 +49,8 @@ func selectConfig() *assets.ClientConfig {
 	err := forms.Select("Select a server:", keys, &selection)
 	if err != nil {
 		fmt.Println(err.Error())
-		return nil
+		return "", nil
 	}
 
-	return configs[selection]
+	return selection, configs[selection]
 }

client/cli/connection_lost.go

@@ -0,0 +1,26 @@
+package cli
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
+)
+
+func handleConnectionLost(ln *grpc.ClientConn) {
+	if ln == nil {
+		return
+	}
+	currentState := ln.GetState()
+	// currentState should be "Ready" when the connection is established.
+	if ln.WaitForStateChange(context.Background(), currentState) {
+		newState := ln.GetState()
+		// newState will be "Idle" if the connection is lost.
+		if newState == connectivity.Idle {
+			fmt.Println("\nLost connection to server. Exiting now.")
+			os.Exit(1)
+		}
+	}
+}

client/cli/console.go

@@ -19,9 +19,7 @@ package cli
 */
 
 import (
-	"context"
 	"fmt"
-	"os"
 
 	"github.com/bishopfox/sliver/client/assets"
 	"github.com/bishopfox/sliver/client/command"
@@ -30,7 +28,6 @@ import (
 	"github.com/bishopfox/sliver/protobuf/rpcpb"
 	"github.com/spf13/cobra"
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/connectivity"
 )
 
 // consoleCmd generates the console with required pre/post runners.
@@ -46,16 +43,14 @@ func consoleCmd(con *console.SliverClient) *cobra.Command {
 }
 
 func consoleRunnerCmd(con *console.SliverClient, run bool) (pre, post func(cmd *cobra.Command, args []string) error) {
-	var ln *grpc.ClientConn
-
 	pre = func(cmd *cobra.Command, _ []string) error {
 
 		configs := assets.GetConfigs()
 		if len(configs) == 0 {
 			fmt.Printf("No config files found at %s (see --help)\n", assets.GetConfigDir())
 			return nil
 		}
-		config := selectConfig()
+		configKey, config := selectConfig()
 		if config == nil {
 			return nil
 		}
@@ -72,40 +67,20 @@ func consoleRunnerCmd(con *console.SliverClient, run bool) (pre, post func(cmd *
 		}
 
 		var rpc rpcpb.SliverRPCClient
+		var ln *grpc.ClientConn
 
 		rpc, ln, err = transport.MTLSConnect(config)
 		if err != nil {
 			fmt.Printf("Connection to server failed %s", err)
 			return nil
 		}
-
-		// Wait for any connection state changes and exit if the connection is lost.
-		go handleConnectionLost(ln)
-
-		return console.StartClient(con, rpc, command.ServerCommands(con, nil), command.SliverCommands(con), run, rcScript)
+		return console.StartClient(con, rpc, ln, &console.ConnectionDetails{ConfigKey: configKey, Config: config}, command.ServerCommands(con, nil), command.SliverCommands(con), run, rcScript)
 	}
 
 	// Close the RPC connection once exiting
 	post = func(_ *cobra.Command, _ []string) error {
-		if ln != nil {
-			return ln.Close()
-		}
-
-		return nil
+		return con.CloseConnection()
 	}
 
 	return pre, post
 }
-
-func handleConnectionLost(ln *grpc.ClientConn) {
-	currentState := ln.GetState()
-	// currentState should be "Ready" when the connection is established.
-	if ln.WaitForStateChange(context.Background(), currentState) {
-		newState := ln.GetState()
-		// newState will be "Idle" if the connection is lost.
-		if newState == connectivity.Idle {
-			fmt.Println("\nLost connection to server. Exiting now.")
-			os.Exit(1)
-		}
-	}
-}

client/command/armory/armory.go

@@ -428,7 +428,7 @@ func PrintArmoryPackages(aliases []*alias.AliasManifest, exts []*extensions.Exte
 
 	tw := table.NewWriter()
 	tw.SetStyle(settings.GetTableStyle(con))
-	tw.SetTitle(console.Bold + "Packages" + console.Normal)
+	tw.SetTitle(console.StyleBold.Render("Packages"))
 
 	urlMargin := 150 // Extra margin needed to show URL column
 
@@ -493,26 +493,26 @@ func PrintArmoryPackages(aliases []*alias.AliasManifest, exts []*extensions.Exte
 
 	rows := []table.Row{}
 	for _, pkg := range entries {
-		color := console.Normal
+		style := console.StyleNormal
 		if extensions.CmdExists(pkg.CommandName, sliverMenu.Command) {
-			color = console.Green
+			style = console.StyleGreen
 		}
 		if con.Settings.SmallTermWidth+urlMargin < width {
 			rows = append(rows, table.Row{
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.ArmoryName),
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.CommandName),
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.Version),
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.Type),
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.Help),
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.URL),
+				style.Render(pkg.ArmoryName),
+				style.Render(pkg.CommandName),
+				style.Render(pkg.Version),
+				style.Render(pkg.Type),
+				style.Render(pkg.Help),
+				style.Render(pkg.URL),
 			})
 		} else {
 			rows = append(rows, table.Row{
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.ArmoryName),
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.CommandName),
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.Version),
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.Type),
-				fmt.Sprintf(color+"%s"+console.Normal, pkg.Help),
+				style.Render(pkg.ArmoryName),
+				style.Render(pkg.CommandName),
+				style.Render(pkg.Version),
+				style.Render(pkg.Type),
+				style.Render(pkg.Help),
 			})
 		}
 	}
@@ -524,7 +524,7 @@ func PrintArmoryPackages(aliases []*alias.AliasManifest, exts []*extensions.Exte
 func PrintArmoryBundles(bundles []*ArmoryBundle, con *console.SliverClient) {
 	tw := table.NewWriter()
 	tw.SetStyle(settings.GetTableStyle(con))
-	tw.SetTitle(console.Bold + "Bundles" + console.Normal)
+	tw.SetTitle(console.StyleBold.Render("Bundles"))
 	tw.AppendHeader(table.Row{
 		"Armory Name",
 		"Name",

client/command/armory/update.go

@@ -257,7 +257,7 @@ func displayAvailableUpdates(con *console.SliverClient, updateKeys []UpdateIdent
 	if len(extensionUpdates) != 1 {
 		extensionSuffix = "s"
 	}
-	tw.SetTitle(console.Bold + fmt.Sprintf(title, len(aliasUpdates), aliasSuffix, len(extensionUpdates), extensionSuffix) + console.Normal)
+	tw.SetTitle(console.StyleBold.Render(fmt.Sprintf(title, len(aliasUpdates), aliasSuffix, len(extensionUpdates), extensionSuffix)))
 	tw.AppendHeader(table.Row{
 		"Package Name",
 		"Package Type",