diff --git a/internal/ui/infer_test.go b/internal/ui/infer_test.go index 6d89e8e..0f3f393 100644 --- a/internal/ui/infer_test.go +++ b/internal/ui/infer_test.go @@ -16,24 +16,24 @@ func TestExtractEmbeddedVars(t *testing.T) { }{ { name: "simple extraction", - template: "-p $credential", - actual: "-p mypassword", + template: "--value $setting", + actual: "--value blue", scope: map[string]string{}, - expected: map[string]string{"credential": "mypassword"}, + expected: map[string]string{"setting": "blue"}, }, { name: "extraction with colon", - template: "-p $credential", - actual: "-p :mypassword", + template: "--value $setting", + actual: "--value :blue", scope: map[string]string{}, - expected: map[string]string{"credential": ":mypassword"}, + expected: map[string]string{"setting": ":blue"}, }, { - name: "hash extraction", - template: "-H $credential", - actual: "-H aad3b435b51404eeaad3b435b51404ee:abc123", + name: "compound value extraction", + template: "--value $setting", + actual: "--value alpha:beta", scope: map[string]string{}, - expected: map[string]string{"credential": "aad3b435b51404eeaad3b435b51404ee:abc123"}, + expected: map[string]string{"setting": "alpha:beta"}, }, } @@ -65,34 +65,34 @@ func TestPrefillScopeFromMatch(t *testing.T) { expectedScope map[string]string }{ { - name: "auth_flags extraction with -k", - command: "bloodyad --host $dc_ip -d $domain $auth_flags", - input: "bloodyad --host 10.0.0.1 -d test.local -k", + name: "mode_flags extraction with dry run", + command: "deployctl --server $server -p $project $mode_flags", + input: "deployctl --server app01 -p website --dry-run", expectedScope: map[string]string{ - "dc_ip": "10.0.0.1", - "domain": "test.local", - "auth_flags": "-k", + "server": "app01", + "project": "website", + "mode_flags": "--dry-run", }, }, { - name: "auth_flags extraction with -p password", - command: "bloodyad --host $dc_ip -d $domain $auth_flags", - input: "bloodyad --host 10.0.0.1 -d test.local -p mypassword", + name: "mode_flags extraction with labeled value", + command: "deployctl --server $server -p $project $mode_flags", + input: "deployctl --server app01 -p website --tag stable", expectedScope: map[string]string{ - "dc_ip": "10.0.0.1", - "domain": "test.local", - "auth_flags": "-p mypassword", + "server": "app01", + "project": "website", + "mode_flags": "--tag stable", }, }, { - name: "full bloodyAD command with mid-command auth_flags", - command: "bloodyAD --host $rhost_name -d $domain -u $user $auth_flags add badSuccessor $rhost_name", - input: "bloodyAD --host test -d bacon.htb -u Administrator -p 123 add badSuccessor test", + name: "full deploy command with mid-command flags", + command: "deployctl --server $server -p $project -u $user $mode_flags apply service $server", + input: "deployctl --server app01 -p website -u alice --tag stable apply service app01", expectedScope: map[string]string{ - "rhost_name": "test", - "domain": "bacon.htb", - "user": "Administrator", - "auth_flags": "-p 123", + "server": "app01", + "project": "website", + "user": "alice", + "mode_flags": "--tag stable", }, }, } @@ -120,17 +120,51 @@ func TestPrefillScopeFromMatch(t *testing.T) { } } +func TestFindMatchingCheatPrefersSpecificCommandOverVarOnly(t *testing.T) { + input := "make deploy SERVICE=api ENV=prod LOG=run-$(date +%Y%m%d-%H%M%S).txt" + varOnly := &parser.Cheat{ + Header: "Single value", + Command: "$item_name", + } + deploy := &parser.Cheat{ + Header: "Deploy service", + Command: "make deploy SERVICE=$service ENV=$env LOG=run-$(date +%Y%m%d-%H%M%S).txt", + } + + got := findMatchingCheat([]*parser.Cheat{varOnly, deploy}, input) + if got != deploy { + t.Fatalf("matched %q, want %q", got.Header, deploy.Header) + } + + prefillScopeFromMatch(got, input) + if got.Scope["service"] != "api" { + t.Fatalf("service = %q, want api", got.Scope["service"]) + } + if got.Scope["env"] != "prod" { + t.Fatalf("env = %q, want prod", got.Scope["env"]) + } +} + +func TestFindMatchingCheatDoesNotUseVarOnlyAsCatchAll(t *testing.T) { + got := findMatchingCheat([]*parser.Cheat{ + {Header: "Single value", Command: "$item_name"}, + }, "tool run thing --flag value") + if got != nil { + t.Fatalf("matched %q, want nil", got.Header) + } +} + func TestInferDependentVars(t *testing.T) { index := &parser.CheatIndex{ Modules: map[string]*parser.Module{ - "bloodyad": { - Name: "bloodyad", + "deploy": { + Name: "deploy", Vars: []parser.VarDef{ - {Name: "auth_method", Shell: "echo -e 'kerberos\npassword\nhash'"}, - {Name: "auth_flags", Literal: "-k", Condition: "$auth_method == kerberos"}, - {Name: "auth_flags", Literal: "-p $credential", Condition: "$auth_method == password"}, - {Name: "auth_flags", Literal: "-H $credential", Condition: "$auth_method == hash"}, - {Name: "credential", Shell: ""}, + {Name: "run_mode", Shell: "printf 'preview\npublish\narchive\n'"}, + {Name: "mode_flags", Literal: "--dry-run", Condition: "$run_mode == preview"}, + {Name: "mode_flags", Literal: "--tag $label", Condition: "$run_mode == publish"}, + {Name: "mode_flags", Literal: "--archive $label", Condition: "$run_mode == archive"}, + {Name: "label", Shell: ""}, }, }, }, @@ -143,23 +177,23 @@ func TestInferDependentVars(t *testing.T) { imports []string }{ { - name: "kerberos flag should infer auth_method", - initialScope: map[string]string{"auth_flags": "-k"}, + name: "preview flag should infer run_mode", + initialScope: map[string]string{"mode_flags": "--dry-run"}, expectedScope: map[string]string{ - "auth_flags": "-k", - "auth_method": "kerberos", + "mode_flags": "--dry-run", + "run_mode": "preview", }, - imports: []string{"bloodyad"}, + imports: []string{"deploy"}, }, { - name: "password flag should infer auth_method and credential", - initialScope: map[string]string{"auth_flags": "-p mypassword"}, + name: "publish flag should infer run_mode and label", + initialScope: map[string]string{"mode_flags": "--tag stable"}, expectedScope: map[string]string{ - "auth_flags": "-p mypassword", - "auth_method": "password", - "credential": "mypassword", + "mode_flags": "--tag stable", + "run_mode": "publish", + "label": "stable", }, - imports: []string{"bloodyad"}, + imports: []string{"deploy"}, }, } diff --git a/internal/ui/match.go b/internal/ui/match.go index b455c9b..fc865d6 100644 --- a/internal/ui/match.go +++ b/internal/ui/match.go @@ -11,20 +11,24 @@ import ( // findMatchingCheat finds a cheat whose command pattern matches the input. // It builds a regex from the cheat command (replacing $var with capture groups) -// and returns the first match. +// and returns the most specific match. A command that is only "$var" matches +// almost anything, so first-match behavior can steal a more exact command. func findMatchingCheat(cheats []*parser.Cheat, input string) *parser.Cheat { input = strings.TrimSpace(input) if input == "" { return nil } + var best *parser.Cheat + bestScore := 0 for _, cheat := range cheats { - pattern, _ := buildMatchPattern(cheat.Command) - if pattern.MatchString(input) { - return cheat + pattern, _, score := buildMatchPatternWithScore(cheat.Command) + if pattern.MatchString(input) && score > bestScore { + best = cheat + bestScore = score } } - return nil + return best } // buildMatchPattern converts a command template to a regex pattern for matching. @@ -36,6 +40,11 @@ func findMatchingCheat(cheats []*parser.Cheat, input string) *parser.Cheat { // The slice may have duplicates (one entry per capture group) because Go // regex doesn't support backreferences. func buildMatchPattern(cmd string) (*regexp.Regexp, []string) { + pattern, varOrder, _ := buildMatchPatternWithScore(cmd) + return pattern, varOrder +} + +func buildMatchPatternWithScore(cmd string) (*regexp.Regexp, []string, int) { var parts []string if config.VarSyntaxAllowsDollar() { parts = append(parts, `\$(\w+)`) @@ -50,6 +59,7 @@ func buildMatchPattern(cmd string) (*regexp.Regexp, []string) { allMatches := varPattern.FindAllStringSubmatchIndex(cmd, -1) var varOrder []string + literalScore := 0 var result strings.Builder result.WriteString(`^\s*`) @@ -68,7 +78,9 @@ func buildMatchPattern(cmd string) (*regexp.Regexp, []string) { } if varStart > lastEnd { - result.WriteString(regexp.QuoteMeta(cmd[lastEnd:varStart])) + literal := cmd[lastEnd:varStart] + literalScore += len(strings.TrimSpace(literal)) + result.WriteString(regexp.QuoteMeta(literal)) } varOrder = append(varOrder, varName) @@ -120,15 +132,17 @@ func buildMatchPattern(cmd string) (*regexp.Regexp, []string) { } if lastEnd < len(cmd) { - result.WriteString(regexp.QuoteMeta(cmd[lastEnd:])) + literal := cmd[lastEnd:] + literalScore += len(strings.TrimSpace(literal)) + result.WriteString(regexp.QuoteMeta(literal)) } result.WriteString(`\s*$`) re, err := regexp.Compile(result.String()) if err != nil { - return regexp.MustCompile(`^$`), nil + return regexp.MustCompile(`^$`), nil, 0 } - return re, varOrder + return re, varOrder, literalScore } // prefillScopeFromMatch extracts variable values from the matched command and diff --git a/internal/ui/match_test.go b/internal/ui/match_test.go index da76c22..89dfe44 100644 --- a/internal/ui/match_test.go +++ b/internal/ui/match_test.go @@ -29,7 +29,7 @@ func TestBuildMatchPattern(t *testing.T) { }, { name: "multiple vars", - cmd: "nmap -p $port $host", + cmd: "tool run --port $port $host", wantVarNames: []string{"port", "host"}, }, { @@ -62,11 +62,11 @@ func TestBuildMatchPattern(t *testing.T) { func TestCheatItemMatchesQuery(t *testing.T) { cheat := &parser.Cheat{ - File: "/cheats/networking/nmap.md", - Header: "Port Scan", - Description: "Scan common ports", - Command: "nmap -sV $target", - Tags: []string{"recon", "pentest"}, + File: "/cheats/projects/deploy.md", + Header: "Deploy Service", + Description: "Deploy a selected service", + Command: "tool deploy $service", + Tags: []string{"release", "service"}, } item := newCheatItem(cheat) @@ -75,15 +75,15 @@ func TestCheatItemMatchesQuery(t *testing.T) { words []string want bool }{ - {"matches folder", []string{"network"}, true}, - {"matches file", []string{"nmap"}, true}, - {"matches header", []string{"port"}, true}, - {"matches description", []string{"common"}, true}, - {"matches command", []string{"nmap"}, true}, - {"matches tag", []string{"pentest"}, true}, - {"matches multiple words", []string{"nmap", "port"}, true}, - {"no match", []string{"gobuster"}, false}, - {"partial multi match fails", []string{"nmap", "gobuster"}, false}, + {"matches folder", []string{"projects"}, true}, + {"matches file", []string{"deploy"}, true}, + {"matches header", []string{"service"}, true}, + {"matches description", []string{"selected"}, true}, + {"matches command", []string{"tool"}, true}, + {"matches tag", []string{"release"}, true}, + {"matches multiple words", []string{"deploy", "service"}, true}, + {"no match", []string{"invoice"}, false}, + {"partial multi match fails", []string{"deploy", "invoice"}, false}, } for _, tt := range tests { diff --git a/internal/ui/path_completion_test.go b/internal/ui/path_completion_test.go index 3d5b24a..d0b08b2 100644 --- a/internal/ui/path_completion_test.go +++ b/internal/ui/path_completion_test.go @@ -46,17 +46,17 @@ func TestCompletePathValueEscapesSpaces(t *testing.T) { func TestCompletePathValueExpandsEnvRoot(t *testing.T) { dir := t.TempDir() - if err := os.Mkdir(filepath.Join(dir, "payloads"), 0o755); err != nil { + if err := os.Mkdir(filepath.Join(dir, "reports"), 0o755); err != nil { t.Fatal(err) } t.Setenv("CHEATMD_COMPLETION_ROOT", dir) - input := "$CHEATMD_COMPLETION_ROOT/pay" + input := "$CHEATMD_COMPLETION_ROOT/rep" result, ok := completePathValue(input, len([]rune(input))) if !ok { t.Fatal("expected completion") } - if result.Value != "$CHEATMD_COMPLETION_ROOT/payloads/" { + if result.Value != "$CHEATMD_COMPLETION_ROOT/reports/" { t.Fatalf("completion = %q", result.Value) } } diff --git a/internal/ui/selector_test.go b/internal/ui/selector_test.go index 509e601..344e1a0 100644 --- a/internal/ui/selector_test.go +++ b/internal/ui/selector_test.go @@ -115,10 +115,10 @@ func TestGetDisplayColumn(t *testing.T) { }, { name: "comma delimited column 1", - line: "admin,password123,active", + line: "primary,blue,active", delimiter: ",", column: 1, - want: "admin", + want: "primary", }, { name: "column out of range returns full line", diff --git a/internal/ui/var_resolve.go b/internal/ui/var_resolve.go index c2ceb9e..40df2c0 100644 --- a/internal/ui/var_resolve.go +++ b/internal/ui/var_resolve.go @@ -309,12 +309,11 @@ func (m *mainModel) handleShellResult(msg shellResultMsg) (tea.Model, tea.Cmd) { } else { m.varState.picker.SetItems(items) } - m.varState.picker.Filter(m.textInput.Value()) - if vs.prefill != "" { m.textInput.SetValue(vs.prefill) m.textInput.CursorEnd() } + m.varState.picker.Filter(m.textInput.Value()) } return m, nil diff --git a/internal/ui/var_resolve_test.go b/internal/ui/var_resolve_test.go index 68a9714..682ed61 100644 --- a/internal/ui/var_resolve_test.go +++ b/internal/ui/var_resolve_test.go @@ -3,6 +3,41 @@ package ui import ( "testing" + "github.com/gubarz/cheatmd/pkg/parser" +) + +func TestVarResolvePrefillFiltersSelectionOptions(t *testing.T) { + m := newMainModel([]*parser.Cheat{{Header: "One"}}, parser.NewCheatIndex(), nil) + m.phase = phaseVarResolve + m.varState = &varResolveState{ + vars: []varState{ + { + def: parser.VarDef{Name: "mode"}, + prefill: "publish", + }, + }, + } + + model, _ := m.handleShellResult(shellResultMsg{ + options: []string{ + "preview\tPreview changes", + "publish\tPublish changes", + "archive\tArchive output", + }, + }) + got := model.(*mainModel) + + if got.textInput.Value() != "publish" { + t.Fatalf("input = %q, want publish", got.textInput.Value()) + } + if got.varState.picker == nil { + t.Fatal("picker is nil") + } + if len(got.varState.picker.Filtered) != 1 { + t.Fatalf("filtered options = %d, want 1", len(got.varState.picker.Filtered)) + } + if got.varState.picker.Filtered[0].Display != "publish\tPublish changes" { + t.Fatalf("filtered option = %q", got.varState.picker.Filtered[0].Display) tea "github.com/charmbracelet/bubbletea" "github.com/gubarz/cheatmd/pkg/parser" )