main
1package cmd
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "mysh/pkg/cache"
8 "mysh/pkg/mythic"
9 "os"
10 "path/filepath"
11 "regexp"
12 "strings"
13 "time"
14
15 "github.com/jedib0t/go-pretty/v6/table"
16 "github.com/spf13/cobra"
17)
18
19type rawMatch struct {
20 taskID int
21 sourceType string // "command" or "output"
22 content string
23}
24
25var (
26 grepCallbackID int
27 grepCaseInsensitive bool
28 grepLimit int
29 grepRegex bool
30 grepInvert bool
31 grepSearchCommands bool
32 grepSearchOutput bool
33 grepCacheOnly bool
34 grepRaw bool
35)
36
37var grepCmd = &cobra.Command{
38 Use: "grep <pattern>",
39 Short: "Search through task outputs and commands",
40 Long: `Search through task response outputs and/or command execution strings for a specific pattern.
41Shows a table of matching tasks with their commands instead of full output.
42This is like grep for Mythic task results - find specific strings, commands, or data across all your task history.
43
44By default, searches both command strings and output. Use --commands-only or --output-only to restrict search scope.`,
45 Args: cobra.ExactArgs(1),
46 RunE: runGrep,
47}
48
49func init() {
50 rootCmd.AddCommand(grepCmd)
51 grepCmd.Flags().IntVarP(&grepCallbackID, "callback", "c", 0, "Search only tasks from specific callback ID")
52 grepCmd.Flags().BoolVarP(&grepCaseInsensitive, "ignore-case", "i", false, "Case insensitive search")
53 grepCmd.Flags().IntVarP(&grepLimit, "limit", "l", 0, "Maximum number of tasks to search (0 = no limit)")
54 grepCmd.Flags().BoolVarP(&grepRegex, "regexp", "E", false, "Treat pattern as regular expression")
55 grepCmd.Flags().BoolVarP(&grepInvert, "invert-match", "v", false, "Show tasks that don't match")
56 grepCmd.Flags().BoolVar(&grepSearchCommands, "commands-only", false, "Search only in command execution strings (command + params)")
57 grepCmd.Flags().BoolVar(&grepSearchOutput, "output-only", false, "Search only in task response outputs")
58 grepCmd.Flags().BoolVar(&grepCacheOnly, "cache-only", false, "Search only cached tasks (don't fetch new data from server)")
59 grepCmd.Flags().BoolVar(&grepRaw, "raw", false, "Display raw output with grep -R style single-line file delimiters")
60}
61
62func runGrep(cmd *cobra.Command, args []string) error {
63 if err := validateConfig(); err != nil {
64 return err
65 }
66
67 // Validate search flags
68 if grepSearchCommands && grepSearchOutput {
69 return fmt.Errorf("cannot use both --commands-only and --output-only flags together")
70 }
71
72 // Set defaults: if neither flag is specified, search both
73 searchCommands := true
74 searchOutput := true
75 if grepSearchCommands {
76 searchOutput = false
77 } else if grepSearchOutput {
78 searchCommands = false
79 }
80
81 pattern := args[0]
82 client := mythic.NewClient(mythicURL, token, insecure, socksProxy)
83
84 // Initialize cache
85 taskCache, err := cache.New(mythicURL)
86 if err != nil {
87 // If cache initialization fails, continue without caching
88 taskCache = nil
89 }
90
91 // Use shorter timeout for getting task lists
92 listCtx, listCancel := context.WithTimeout(context.Background(), 30*time.Second)
93 defer listCancel()
94
95 // Determine search scope description
96 var searchScope string
97 if searchCommands && searchOutput {
98 searchScope = "commands and outputs"
99 } else if searchCommands {
100 searchScope = "command strings"
101 } else {
102 searchScope = "outputs"
103 }
104
105 var tasks []mythic.Task
106
107 if grepCacheOnly {
108 // Cache-only mode: only search cached tasks
109 fmt.Printf("Searching %s in cached tasks only for pattern: %s\n", searchScope, pattern)
110 tasks = getCachedTasks(taskCache)
111 if len(tasks) == 0 {
112 fmt.Println("No cached tasks found to search")
113 return nil
114 }
115 fmt.Printf("Searching through %d cached tasks...\n", len(tasks))
116 } else {
117 // Normal mode: get fresh data from server
118 // Step 1: Get lightweight task metadata first
119 if grepCallbackID > 0 {
120 tasks, err = client.GetCallbackTasks(listCtx, grepCallbackID)
121 if err != nil {
122 return fmt.Errorf("failed to get tasks for callback %d: %w", grepCallbackID, err)
123 }
124 fmt.Printf("Searching %s from callback %d for pattern: %s\n", searchScope, grepCallbackID, pattern)
125 } else {
126 tasks, err = client.GetAllTasks(listCtx, grepLimit)
127 if err != nil {
128 return fmt.Errorf("failed to get all tasks: %w", err)
129 }
130 fmt.Printf("Searching %s in all tasks for pattern: %s\n", searchScope, pattern)
131 }
132
133 if len(tasks) == 0 {
134 fmt.Println("No tasks found to search")
135 return nil
136 }
137
138 fmt.Printf("Searching through %d tasks...\n", len(tasks))
139
140 // Step 2: For tasks that need response data, populate from cache or fetch
141 if searchOutput {
142 fmt.Printf("Loading response data (using cache when available)...\n")
143 tasks = populateTaskResponses(listCtx, client, tasks, taskCache)
144 }
145 }
146
147 fmt.Println()
148
149 // Prepare pattern for matching
150 var matcher func(string) bool
151 if grepRegex {
152 regexPattern := pattern
153 if grepCaseInsensitive {
154 regexPattern = "(?i)" + pattern
155 }
156 re, err := regexp.Compile(regexPattern)
157 if err != nil {
158 return fmt.Errorf("invalid regular expression: %w", err)
159 }
160 matcher = re.MatchString
161 } else {
162 searchPattern := pattern
163 if grepCaseInsensitive {
164 searchPattern = strings.ToLower(pattern)
165 }
166 matcher = func(text string) bool {
167 searchText := text
168 if grepCaseInsensitive {
169 searchText = strings.ToLower(text)
170 }
171 return strings.Contains(searchText, searchPattern)
172 }
173 }
174
175 // Search through tasks in memory (much faster since we have all data)
176 var matchingTasks []mythic.Task
177 var rawMatches []rawMatch
178 processedTasks := 0
179
180 for _, task := range tasks {
181 processedTasks++
182 var matches bool
183 var taskMatches []rawMatch
184
185 // Check command string if enabled
186 if searchCommands {
187 commandString := task.Command
188 if task.Params != "" {
189 commandString += " " + task.Params
190 }
191 if matcher(commandString) {
192 matches = true
193 if grepRaw {
194 taskMatches = append(taskMatches, rawMatch{
195 taskID: task.DisplayID,
196 sourceType: "command",
197 content: commandString,
198 })
199 }
200 }
201 }
202
203 // Check response output if enabled
204 if searchOutput && task.Response != "" {
205 if matcher(task.Response) {
206 matches = true
207 if grepRaw {
208 taskMatches = append(taskMatches, rawMatch{
209 taskID: task.DisplayID,
210 sourceType: "output",
211 content: task.Response,
212 })
213 }
214 }
215 }
216
217 // Apply invert logic
218 if grepInvert {
219 matches = !matches
220 }
221
222 if matches {
223 matchingTasks = append(matchingTasks, task)
224 if grepRaw {
225 rawMatches = append(rawMatches, taskMatches...)
226 }
227 }
228 }
229
230 if len(matchingTasks) == 0 {
231 fmt.Printf("No matches found in %d tasks searched\n", processedTasks)
232 return nil
233 }
234
235 // Show results
236 if grepRaw {
237 showRawMatches(rawMatches)
238 } else {
239 showMatchingTasksTable(matchingTasks)
240 }
241
242 fmt.Printf("\nFound %d matches in %d tasks searched\n", len(matchingTasks), processedTasks)
243 return nil
244}
245
246func showRawMatches(matches []rawMatch) {
247 for i, match := range matches {
248 // Format like tail -f: ==> filename <==
249 // Use task ID and source type as the "filename"
250 filename := fmt.Sprintf("task_%d_%s", match.taskID, match.sourceType)
251
252 // Add separator between files (except for the first one)
253 if i > 0 {
254 fmt.Println()
255 }
256
257 fmt.Printf("==> %s <==\n", filename)
258 fmt.Print(match.content)
259
260 // Ensure content ends with newline if it doesn't already
261 if !strings.HasSuffix(match.content, "\n") {
262 fmt.Println()
263 }
264 }
265}
266
267func showMatchingTasksTable(tasks []mythic.Task) {
268 // Create a table for formatted output
269 t := table.NewWriter()
270 t.SetOutputMirror(os.Stdout)
271
272 // Configure table styling - header borders only
273 t.SetStyle(table.Style{
274 Name: "HeaderOnlyBorders",
275 Box: table.StyleBoxDefault,
276 Options: table.Options{
277 DrawBorder: false, // No outer border
278 SeparateColumns: false, // No column separators
279 SeparateFooter: false, // No footer separator
280 SeparateHeader: true, // Keep header separator
281 SeparateRows: false, // No row separators
282 },
283 })
284
285 // Define columns with priorities (1=highest, higher numbers can be dropped)
286 columns := []ColumnInfo{
287 {Header: "TASK ID", Number: 1, MinWidth: 6, MaxWidth: 8, Priority: 1}, // Always show
288 {Header: "COMMAND", Number: 2, MinWidth: 8, MaxWidth: 12, Priority: 2}, // High priority
289 {Header: "PARAMS", Number: 3, MinWidth: 10, MaxWidth: 30, Priority: 5}, // Can drop
290 {Header: "STATUS", Number: 4, MinWidth: 6, MaxWidth: 10, Priority: 3}, // Important
291 {Header: "CALLBACK", Number: 5, MinWidth: 6, MaxWidth: 8, Priority: 4}, // Important
292 {Header: "TIMESTAMP", Number: 6, MinWidth: 10, MaxWidth: 14, Priority: 6}, // First to drop
293 }
294
295 // Configure table for current terminal width
296 dt := configureTableForTerminal(t, columns)
297
298 for _, task := range tasks {
299 // Use full params - table library will handle truncation
300 params := task.Params
301
302 // Format timestamp
303 timestamp := task.Timestamp
304 if timestamp != "" {
305 if t, err := time.Parse(time.RFC3339, timestamp); err == nil {
306 timestamp = t.Format("2006-01-02 15:04")
307 }
308 }
309
310 // Prepare all column data
311 allData := []interface{}{
312 task.DisplayID,
313 task.Command,
314 params,
315 task.Status,
316 task.CallbackID,
317 timestamp,
318 }
319
320 // Add row with only visible columns
321 dt.AppendRowForColumns(allData)
322 }
323
324 t.Render()
325}
326
327// populateTaskResponses efficiently loads response data for tasks using cache when possible
328func populateTaskResponses(ctx context.Context, client *mythic.Client, tasks []mythic.Task, taskCache *cache.TaskCache) []mythic.Task {
329 var tasksWithResponses []mythic.Task
330 var tasksNeedingFetch []mythic.Task
331 cacheHits := 0
332 skippedIncomplete := 0
333
334 // Step 1: Check cache for completed tasks, skip incomplete tasks
335 for _, task := range tasks {
336 // Skip incomplete tasks - they shouldn't be searched for responses
337 if !task.Completed && task.Status != "completed" && task.Status != "error" {
338 skippedIncomplete++
339 // Add task with empty response for command-only searches
340 tasksWithResponses = append(tasksWithResponses, task)
341 continue
342 }
343
344 // Try to get from cache
345 var foundInCache bool
346 if taskCache != nil {
347 if cachedTask, found := taskCache.GetCachedTask(task.ID, mythicURL); found {
348 // Use cached task with response data
349 tasksWithResponses = append(tasksWithResponses, *cachedTask)
350 foundInCache = true
351 cacheHits++
352 }
353 }
354
355 if !foundInCache {
356 // Need to fetch this task's response
357 tasksNeedingFetch = append(tasksNeedingFetch, task)
358 }
359 }
360
361 // Step 2: Batch fetch missing responses
362 if len(tasksNeedingFetch) > 0 {
363 fmt.Printf("Cache hits: %d, fetching: %d, skipped incomplete: %d\n",
364 cacheHits, len(tasksNeedingFetch), skippedIncomplete)
365
366 // Extract task IDs for batch query
367 taskIDs := make([]int, len(tasksNeedingFetch))
368 for i, task := range tasksNeedingFetch {
369 taskIDs[i] = task.ID
370 }
371
372 // Batch fetch all missing task responses in one query
373 fetchedTasks, err := client.GetTasksWithResponsesByIDs(ctx, taskIDs)
374 if err != nil {
375 fmt.Fprintf(os.Stderr, "Warning: batch fetch failed, falling back to individual requests: %v\n", err)
376 // Fallback to individual requests if batch fails
377 for _, task := range tasksNeedingFetch {
378 fullTask, err := client.GetTaskResponse(ctx, task.ID)
379 if err != nil {
380 fmt.Fprintf(os.Stderr, "Warning: failed to get response for task %d: %v\n", task.DisplayID, err)
381 tasksWithResponses = append(tasksWithResponses, task)
382 continue
383 }
384 tasksWithResponses = append(tasksWithResponses, *fullTask)
385 }
386 } else {
387 // Successfully batch fetched - add all to results
388 for _, fetchedTask := range fetchedTasks {
389 tasksWithResponses = append(tasksWithResponses, fetchedTask)
390 }
391
392 // Cache the batch results
393 cached := 0
394 notCached := 0
395 if taskCache != nil {
396 for _, fetchedTask := range fetchedTasks {
397 if err := taskCache.CacheTask(&fetchedTask, mythicURL); err != nil {
398 fmt.Fprintf(os.Stderr, "Warning: failed to cache task response: %v\n", err)
399 } else {
400 // Check if it was actually cached (completed) or skipped (incomplete)
401 if fetchedTask.Completed || fetchedTask.Status == "completed" || fetchedTask.Status == "error" {
402 cached++
403 } else {
404 notCached++
405 }
406 }
407 }
408 if cached > 0 || notCached > 0 {
409 fmt.Printf("Newly cached: %d, not cached (incomplete): %d\n", cached, notCached)
410 }
411 }
412 }
413 } else if cacheHits > 0 || skippedIncomplete > 0 {
414 fmt.Printf("Cache hits: %d, fetching: %d, skipped incomplete: %d\n",
415 cacheHits, 0, skippedIncomplete)
416 }
417
418 return tasksWithResponses
419}
420
421// getCachedTasks loads all cached tasks from the cache directory
422func getCachedTasks(taskCache *cache.TaskCache) []mythic.Task {
423 if taskCache == nil {
424 return []mythic.Task{}
425 }
426
427 // Get all cached task files
428 cacheDir, _, _, err := taskCache.GetCacheInfo()
429 if err != nil {
430 return []mythic.Task{}
431 }
432
433 entries, err := os.ReadDir(cacheDir)
434 if err != nil {
435 return []mythic.Task{}
436 }
437
438 var tasks []mythic.Task
439 for _, entry := range entries {
440 if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
441 continue
442 }
443
444 cachePath := filepath.Join(cacheDir, entry.Name())
445 file, err := os.Open(cachePath)
446 if err != nil {
447 continue
448 }
449
450 var cachedTask cache.CachedTask
451 if err := json.NewDecoder(file).Decode(&cachedTask); err != nil {
452 file.Close()
453 continue
454 }
455 file.Close()
456
457 if cachedTask.Task != nil {
458 tasks = append(tasks, *cachedTask.Task)
459 }
460 }
461
462 return tasks
463}