feat(console): improve context-aware CLI completion

Author: 404tk

runner/console/completer.go

@@ -9,22 +9,6 @@ import (
 	"github.com/404tk/go-prompt"
 )
 
-var core = []prompt.Suggest{
-	{Text: "help", Description: "help menu"},
-	{Text: "use", Description: "use provider"},
-	{Text: "sessions", Description: "list cache credential"},
-	{Text: "note", Description: "add remarks to the session"},
-	{Text: "clear", Description: "clear screen"},
-	{Text: "exit", Description: "exit console"},
-}
-
-var action = []prompt.Suggest{
-	{Text: "show", Description: "show options or payloads"},
-	{Text: "set", Description: "set option or parameter"},
-	{Text: "run", Description: "run selected payload"},
-	{Text: "shell", Description: "open instance-cmd-check session"},
-}
-
 var modules = func() (m []prompt.Suggest) {
 	for k, v := range plugins.Providers {
 		m = append(m, prompt.Suggest{Text: k, Description: v.Desc()})
@@ -54,10 +38,31 @@ var opt = []prompt.Suggest{
 }
 
 func Complete(d prompt.Document) []prompt.Suggest {
+	return completeForMode(d, HelpModeRoot)
+}
+
+func actionCompleter(d prompt.Document) []prompt.Suggest {
+	return completeForMode(d, HelpModeProvider)
+}
+
+func completeForMode(d prompt.Document, mode HelpMode) []prompt.Suggest {
+	ctx := completionContextForMode(mode)
 	args := completionArgs(d)
 	word := d.GetWordBeforeCursor()
+
+	switch mode {
+	case HelpModeProvider:
+		return providerSuggestions(ctx, args, word)
+	case HelpModeShell:
+		return shellSuggestions(ctx, args, word)
+	default:
+		return rootSuggestions(ctx, args, word)
+	}
+}
+
+func rootSuggestions(_ CompletionContext, args []string, word string) []prompt.Suggest {
 	if len(args) <= 1 {
-		return prompt.FilterHasPrefix(core, word, true)
+		return prompt.FilterHasPrefix(rootCommandSuggestions, word, true)
 	}
 	switch args[0] {
 	case "help":
@@ -66,15 +71,17 @@ func Complete(d prompt.Document) []prompt.Suggest {
 		if len(args) == 2 {
 			return prompt.FilterContains(modules, word, true)
 		}
+	case "sessions":
+		return sessionsSuggestions(args, word)
+	case "note":
+		return noteSuggestions(args, word)
 	}
 	return []prompt.Suggest{}
 }
 
-func actionCompleter(d prompt.Document) []prompt.Suggest {
-	args := completionArgs(d)
-	word := d.GetWordBeforeCursor()
+func providerSuggestions(ctx CompletionContext, args []string, word string) []prompt.Suggest {
 	if len(args) <= 1 {
-		return prompt.FilterHasPrefix(append(core, action...), word, true)
+		return prompt.FilterHasPrefix(providerCommandSuggestions(), word, true)
 	}
 	switch args[0] {
 	case "help":
@@ -84,23 +91,62 @@ func actionCompleter(d prompt.Document) []prompt.Suggest {
 			return prompt.FilterContains(modules, word, true)
 		}
 	case "show":
-		if len(args) == 2 {
-			return prompt.FilterContains(opt, word, true)
-		}
+		return showSuggestions(args, word)
 	case "set":
-		if len(args) == 2 {
-			return prompt.FilterContains(optionSuggestions(), word, true)
-		}
-		if len(args) == 3 && args[1] == utils.Payload {
-			return prompt.FilterContains(payloadSuggestions(), word, true)
-		}
-		if len(args) == 3 && args[1] == utils.Version {
-			var versions = []prompt.Suggest{
-				{Text: "Intl", Description: "International Edition"},
-				{Text: "China", Description: "Chinese Edition"},
-			}
-			return prompt.FilterContains(versions, word, true)
+		return setSuggestions(ctx, args, word)
+	case "shell":
+		return shellTargetSuggestions(ctx, args, word)
+	case "sessions":
+		return sessionsSuggestions(args, word)
+	case "note":
+		return noteSuggestions(args, word)
+	}
+	return []prompt.Suggest{}
+}
+
+func shellSuggestions(_ CompletionContext, args []string, word string) []prompt.Suggest {
+	if len(args) <= 1 {
+		return prompt.FilterHasPrefix(shellCommandSuggestions, word, true)
+	}
+	switch args[0] {
+	case "help":
+		return helpSuggestions(args, word)
+	}
+	return []prompt.Suggest{}
+}
+
+func showSuggestions(args []string, word string) []prompt.Suggest {
+	if len(args) == 2 {
+		return prompt.FilterContains(showTopicSuggestions(), word, true)
+	}
+	return []prompt.Suggest{}
+}
+
+func setSuggestions(ctx CompletionContext, args []string, word string) []prompt.Suggest {
+	if len(args) == 2 {
+		return prompt.FilterContains(optionSuggestions(ctx), word, true)
+	}
+	if len(args) != 3 {
+		return []prompt.Suggest{}
+	}
+	switch args[1] {
+	case utils.Payload:
+		return prompt.FilterContains(payloadSuggestions(), word, true)
+	case utils.Version:
+		if _, ok := ctx.Config[utils.Version]; ok {
+			return prompt.FilterContains(versionSuggestionsData, word, true)
 		}
+	case utils.Region:
+		return prompt.FilterContains(getProviderRegionSuggestions(ctx.Provider), word, true)
+	case utils.Metadata:
+		return prompt.FilterContains(getPayloadMetadataSuggestions(ctx.Payload), word, true)
+	}
+	return []prompt.Suggest{}
+}
+
+func shellTargetSuggestions(ctx CompletionContext, args []string, word string) []prompt.Suggest {
+	if len(args) == 2 {
+		return prompt.FilterContains(getShellTargetSuggestions(ctx), word, true)
 	}
 	return []prompt.Suggest{}
 }
@@ -137,9 +183,33 @@ func helpTopicSuggestions() []prompt.Suggest {
 	return suggestions
 }
 
-func optionSuggestions() []prompt.Suggest {
-	suggestions := make([]prompt.Suggest, 0, len(config))
-	for k := range config {
+func optionSuggestions(ctx CompletionContext) []prompt.Suggest {
+	keys := []string{utils.Payload, utils.Metadata}
+	seen := map[string]struct{}{
+		utils.Payload:  {},
+		utils.Metadata: {},
+	}
+
+	for _, key := range providerConfigKeys() {
+		if _, ok := seen[key]; ok {
+			continue
+		}
+		keys = append(keys, key)
+		seen[key] = struct{}{}
+	}
+	for _, key := range []string{utils.Region, utils.Version} {
+		if _, ok := ctx.Config[key]; !ok {
+			continue
+		}
+		if _, ok := seen[key]; ok {
+			continue
+		}
+		keys = append(keys, key)
+		seen[key] = struct{}{}
+	}
+
+	suggestions := make([]prompt.Suggest, 0, len(keys))
+	for _, k := range keys {
 		if v, ok := optionsDesc[k]; ok {
 			suggestions = append(suggestions, prompt.Suggest{Text: k, Description: v})
 		}