Commit 75a52f4
Changed files (8)
pkg
cache
mythic
cmd/cache.go
@@ -0,0 +1,130 @@
+package cmd
+
+import (
+ "context"
+ "fmt"
+ "mysh/pkg/cache"
+ "mysh/pkg/mythic"
+ "time"
+
+ "github.com/spf13/cobra"
+)
+
+var cacheCmd = &cobra.Command{
+ Use: "cache",
+ Short: "Show cache information",
+ Long: "Display information about the task result cache including location, size, and number of cached tasks.",
+ RunE: runCacheInfo,
+}
+
+var cacheFlushCmd = &cobra.Command{
+ Use: "flush",
+ Short: "Flush all cached task results",
+ Long: "Remove all cached task results from the cache directory.",
+ RunE: runCacheFlush,
+}
+
+var cacheUpdateCmd = &cobra.Command{
+ Use: "update",
+ Short: "Update cache with latest task results",
+ Long: "Fetch all tasks from the server and update the local cache with completed task results.",
+ RunE: runCacheUpdate,
+}
+
+func init() {
+ rootCmd.AddCommand(cacheCmd)
+ cacheCmd.AddCommand(cacheFlushCmd)
+ cacheCmd.AddCommand(cacheUpdateCmd)
+}
+
+func runCacheInfo(cmd *cobra.Command, args []string) error {
+ if err := validateConfig(); err != nil {
+ return err
+ }
+
+ taskCache, err := cache.New(mythicURL)
+ if err != nil {
+ return fmt.Errorf("failed to initialize cache: %w", err)
+ }
+
+ cacheDir, fileCount, totalSize, err := taskCache.GetCacheInfo()
+ if err != nil {
+ return fmt.Errorf("failed to get cache info: %w", err)
+ }
+
+ fmt.Printf("Cache Information:\n")
+ fmt.Printf("Location: %s\n", cacheDir)
+ fmt.Printf("Cached tasks: %d\n", fileCount)
+ fmt.Printf("Total size: %s\n", formatBytes(totalSize))
+
+ return nil
+}
+
+func runCacheFlush(cmd *cobra.Command, args []string) error {
+ if err := validateConfig(); err != nil {
+ return err
+ }
+
+ taskCache, err := cache.New(mythicURL)
+ if err != nil {
+ return fmt.Errorf("failed to initialize cache: %w", err)
+ }
+
+ if err := taskCache.ClearCache(); err != nil {
+ return fmt.Errorf("failed to flush cache: %w", err)
+ }
+
+ fmt.Println("Cache flushed successfully")
+ return nil
+}
+
+func runCacheUpdate(cmd *cobra.Command, args []string) error {
+ if err := validateConfig(); err != nil {
+ return err
+ }
+
+ client := mythic.NewClient(mythicURL, token, insecure, socksProxy)
+ ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+ defer cancel()
+
+ // Initialize cache
+ taskCache, err := cache.New(mythicURL)
+ if err != nil {
+ return fmt.Errorf("failed to initialize cache: %w", err)
+ }
+
+ fmt.Println("Fetching all tasks from server...")
+
+ // Get all tasks (lightweight metadata first)
+ tasks, err := client.GetAllTasks(ctx, 0) // 0 = no limit
+ if err != nil {
+ return fmt.Errorf("failed to get tasks: %w", err)
+ }
+
+ if len(tasks) == 0 {
+ fmt.Println("No tasks found on server")
+ return nil
+ }
+
+ fmt.Printf("Found %d tasks, updating cache...\n", len(tasks))
+
+ // Use the same logic as grep to populate task responses and update cache
+ updatedTasks := populateTaskResponses(ctx, client, tasks, taskCache)
+
+ fmt.Printf("Cache update completed. Processed %d tasks.\n", len(updatedTasks))
+ return nil
+}
+
+// formatBytes formats byte count as human-readable string
+func formatBytes(bytes int64) string {
+ const unit = 1024
+ if bytes < unit {
+ return fmt.Sprintf("%d B", bytes)
+ }
+ div, exp := int64(unit), 0
+ for n := bytes / unit; n >= unit; n /= unit {
+ div *= unit
+ exp++
+ }
+ return fmt.Sprintf("%.1f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
+}
\ No newline at end of file
cmd/grep.go
@@ -2,9 +2,12 @@ package cmd
import (
"context"
+ "encoding/json"
"fmt"
+ "mysh/pkg/cache"
"mysh/pkg/mythic"
"os"
+ "path/filepath"
"regexp"
"strings"
"text/tabwriter"
@@ -21,6 +24,7 @@ var (
grepInvert bool
grepSearchCommands bool
grepSearchOutput bool
+ grepCacheOnly bool
)
var grepCmd = &cobra.Command{
@@ -44,6 +48,7 @@ func init() {
grepCmd.Flags().BoolVarP(&grepInvert, "invert-match", "v", false, "Show tasks that don't match")
grepCmd.Flags().BoolVar(&grepSearchCommands, "commands-only", false, "Search only in command execution strings (command + params)")
grepCmd.Flags().BoolVar(&grepSearchOutput, "output-only", false, "Search only in task response outputs")
+ grepCmd.Flags().BoolVar(&grepCacheOnly, "cache-only", false, "Search only cached tasks (don't fetch new data from server)")
}
@@ -69,13 +74,17 @@ func runGrep(cmd *cobra.Command, args []string) error {
pattern := args[0]
client := mythic.NewClient(mythicURL, token, insecure, socksProxy)
+ // Initialize cache
+ taskCache, err := cache.New(mythicURL)
+ if err != nil {
+ // If cache initialization fails, continue without caching
+ taskCache = nil
+ }
+
// Use shorter timeout for getting task lists
listCtx, listCancel := context.WithTimeout(context.Background(), 30*time.Second)
defer listCancel()
- var tasks []mythic.Task
- var err error
-
// Determine search scope description
var searchScope string
if searchCommands && searchOutput {
@@ -86,28 +95,49 @@ func runGrep(cmd *cobra.Command, args []string) error {
searchScope = "outputs"
}
- if grepCallbackID > 0 {
- // Search tasks for specific callback - use efficient query that includes responses
- tasks, err = client.GetTasksWithResponses(listCtx, grepCallbackID, grepLimit)
- if err != nil {
- return fmt.Errorf("failed to get tasks for callback %d: %w", grepCallbackID, err)
+ var tasks []mythic.Task
+
+ if grepCacheOnly {
+ // Cache-only mode: only search cached tasks
+ fmt.Printf("Searching %s in cached tasks only for pattern: %s\n", searchScope, pattern)
+ tasks = getCachedTasks(taskCache)
+ if len(tasks) == 0 {
+ fmt.Println("No cached tasks found to search")
+ return nil
}
- fmt.Printf("Searching %s from callback %d for pattern: %s\n", searchScope, grepCallbackID, pattern)
+ fmt.Printf("Searching through %d cached tasks...\n", len(tasks))
} else {
- // Get all tasks from all callbacks with responses in a single query
- tasks, err = client.GetAllTasksWithResponses(listCtx, grepLimit)
- if err != nil {
- return fmt.Errorf("failed to get all tasks with responses: %w", err)
+ // Normal mode: get fresh data from server
+ // Step 1: Get lightweight task metadata first
+ if grepCallbackID > 0 {
+ tasks, err = client.GetCallbackTasks(listCtx, grepCallbackID)
+ if err != nil {
+ return fmt.Errorf("failed to get tasks for callback %d: %w", grepCallbackID, err)
+ }
+ fmt.Printf("Searching %s from callback %d for pattern: %s\n", searchScope, grepCallbackID, pattern)
+ } else {
+ tasks, err = client.GetAllTasks(listCtx, grepLimit)
+ if err != nil {
+ return fmt.Errorf("failed to get all tasks: %w", err)
+ }
+ fmt.Printf("Searching %s in all tasks for pattern: %s\n", searchScope, pattern)
}
- fmt.Printf("Searching %s in all tasks for pattern: %s\n", searchScope, pattern)
- }
- if len(tasks) == 0 {
- fmt.Println("No tasks found to search")
- return nil
+ if len(tasks) == 0 {
+ fmt.Println("No tasks found to search")
+ return nil
+ }
+
+ fmt.Printf("Searching through %d tasks...\n", len(tasks))
+
+ // Step 2: For tasks that need response data, populate from cache or fetch
+ if searchOutput {
+ fmt.Printf("Loading response data (using cache when available)...\n")
+ tasks = populateTaskResponses(listCtx, client, tasks, taskCache)
+ }
}
- fmt.Printf("Searching through %d tasks...\n\n", len(tasks))
+ fmt.Println()
// Prepare pattern for matching
var matcher func(string) bool
@@ -212,4 +242,142 @@ func showMatchingTasksTable(tasks []mythic.Task) {
}
w.Flush()
+}
+
+// populateTaskResponses efficiently loads response data for tasks using cache when possible
+func populateTaskResponses(ctx context.Context, client *mythic.Client, tasks []mythic.Task, taskCache *cache.TaskCache) []mythic.Task {
+ var tasksWithResponses []mythic.Task
+ var tasksNeedingFetch []mythic.Task
+ cacheHits := 0
+ skippedIncomplete := 0
+
+ // Step 1: Check cache for completed tasks, skip incomplete tasks
+ for _, task := range tasks {
+ // Skip incomplete tasks - they shouldn't be searched for responses
+ if !task.Completed && task.Status != "completed" && task.Status != "error" {
+ skippedIncomplete++
+ // Add task with empty response for command-only searches
+ tasksWithResponses = append(tasksWithResponses, task)
+ continue
+ }
+
+ // Try to get from cache
+ var foundInCache bool
+ if taskCache != nil {
+ if cachedTask, found := taskCache.GetCachedTask(task.ID, mythicURL); found {
+ // Use cached task with response data
+ tasksWithResponses = append(tasksWithResponses, *cachedTask)
+ foundInCache = true
+ cacheHits++
+ }
+ }
+
+ if !foundInCache {
+ // Need to fetch this task's response
+ tasksNeedingFetch = append(tasksNeedingFetch, task)
+ }
+ }
+
+ // Step 2: Batch fetch missing responses
+ if len(tasksNeedingFetch) > 0 {
+ fmt.Printf("Cache hits: %d, fetching: %d, skipped incomplete: %d\n",
+ cacheHits, len(tasksNeedingFetch), skippedIncomplete)
+
+ // Extract task IDs for batch query
+ taskIDs := make([]int, len(tasksNeedingFetch))
+ for i, task := range tasksNeedingFetch {
+ taskIDs[i] = task.ID
+ }
+
+ // Batch fetch all missing task responses in one query
+ fetchedTasks, err := client.GetTasksWithResponsesByIDs(ctx, taskIDs)
+ if err != nil {
+ fmt.Fprintf(os.Stderr, "Warning: batch fetch failed, falling back to individual requests: %v\n", err)
+ // Fallback to individual requests if batch fails
+ for _, task := range tasksNeedingFetch {
+ fullTask, err := client.GetTaskResponse(ctx, task.ID)
+ if err != nil {
+ fmt.Fprintf(os.Stderr, "Warning: failed to get response for task %d: %v\n", task.DisplayID, err)
+ tasksWithResponses = append(tasksWithResponses, task)
+ continue
+ }
+ tasksWithResponses = append(tasksWithResponses, *fullTask)
+ }
+ } else {
+ // Successfully batch fetched - add all to results
+ for _, fetchedTask := range fetchedTasks {
+ tasksWithResponses = append(tasksWithResponses, fetchedTask)
+ }
+
+ // Cache the batch results
+ cached := 0
+ notCached := 0
+ if taskCache != nil {
+ for _, fetchedTask := range fetchedTasks {
+ if err := taskCache.CacheTask(&fetchedTask, mythicURL); err != nil {
+ fmt.Fprintf(os.Stderr, "Warning: failed to cache task response: %v\n", err)
+ } else {
+ // Check if it was actually cached (completed) or skipped (incomplete)
+ if fetchedTask.Completed || fetchedTask.Status == "completed" || fetchedTask.Status == "error" {
+ cached++
+ } else {
+ notCached++
+ }
+ }
+ }
+ if cached > 0 || notCached > 0 {
+ fmt.Printf("Newly cached: %d, not cached (incomplete): %d\n", cached, notCached)
+ }
+ }
+ }
+ } else if cacheHits > 0 || skippedIncomplete > 0 {
+ fmt.Printf("Cache hits: %d, fetching: %d, skipped incomplete: %d\n",
+ cacheHits, 0, skippedIncomplete)
+ }
+
+ return tasksWithResponses
+}
+
+// getCachedTasks loads all cached tasks from the cache directory
+func getCachedTasks(taskCache *cache.TaskCache) []mythic.Task {
+ if taskCache == nil {
+ return []mythic.Task{}
+ }
+
+ // Get all cached task files
+ cacheDir, _, _, err := taskCache.GetCacheInfo()
+ if err != nil {
+ return []mythic.Task{}
+ }
+
+ entries, err := os.ReadDir(cacheDir)
+ if err != nil {
+ return []mythic.Task{}
+ }
+
+ var tasks []mythic.Task
+ for _, entry := range entries {
+ if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
+ continue
+ }
+
+ cachePath := filepath.Join(cacheDir, entry.Name())
+ file, err := os.Open(cachePath)
+ if err != nil {
+ continue
+ }
+
+ var cachedTask cache.CachedTask
+ if err := json.NewDecoder(file).Decode(&cachedTask); err != nil {
+ file.Close()
+ continue
+ }
+ file.Close()
+
+ if cachedTask.Task != nil {
+ tasks = append(tasks, *cachedTask.Task)
+ }
+ }
+
+ return tasks
}
\ No newline at end of file
cmd/task_list.go
@@ -6,7 +6,6 @@ import (
"mysh/pkg/mythic"
"os"
"strconv"
- "strings"
"text/tabwriter"
"time"
@@ -14,9 +13,8 @@ import (
)
var (
- listShowResponses bool
- listTaskLimit int
- listShowAll bool
+ listTaskLimit int
+ listShowAll bool
)
var taskListCmd = &cobra.Command{
@@ -30,7 +28,6 @@ var taskListCmd = &cobra.Command{
func init() {
rootCmd.AddCommand(taskListCmd)
- taskListCmd.Flags().BoolVar(&listShowResponses, "responses", false, "Show task responses/output")
taskListCmd.Flags().IntVar(&listTaskLimit, "limit", 10, "Maximum number of tasks to display")
taskListCmd.Flags().BoolVarP(&listShowAll, "all", "a", false, "Show all tasks (same as --limit 0)")
}
@@ -45,7 +42,6 @@ func runTaskList(cmd *cobra.Command, args []string) error {
defer cancel()
var tasks []mythic.Task
- var err error
var showingAllCallbacks bool
// Determine if we're showing tasks for a specific callback or all callbacks
@@ -59,7 +55,8 @@ func runTaskList(cmd *cobra.Command, args []string) error {
limit = 0
}
- tasks, err = client.GetAllTasksWithResponses(ctx, limit)
+ var err error
+ tasks, err = client.GetAllTasks(ctx, limit)
if err != nil {
return fmt.Errorf("failed to get all tasks: %w", err)
}
@@ -70,18 +67,11 @@ func runTaskList(cmd *cobra.Command, args []string) error {
return fmt.Errorf("invalid callback ID: %s", args[0])
}
+ var err error
tasks, err = client.GetCallbackTasks(ctx, callbackID)
if err != nil {
return fmt.Errorf("failed to get tasks for callback %d: %w", callbackID, err)
}
-
- // Apply limit from the end (keep newest tasks if limiting)
- if listShowAll {
- listTaskLimit = 0
- }
- if listTaskLimit > 0 && len(tasks) > listTaskLimit {
- tasks = tasks[len(tasks)-listTaskLimit:]
- }
}
if len(tasks) == 0 {
@@ -95,7 +85,7 @@ func runTaskList(cmd *cobra.Command, args []string) error {
// Display header
if showingAllCallbacks {
- fmt.Printf("All tasks (showing %d, newest first):\n\n", len(tasks))
+ fmt.Printf("All tasks (showing %d, newest last):\n\n", len(tasks))
} else {
fmt.Printf("Tasks for callback %s (showing %d, newest last):\n\n", args[0], len(tasks))
}
@@ -148,45 +138,5 @@ func runTaskList(cmd *cobra.Command, args []string) error {
w.Flush()
- // If responses are requested, fetch and display them
- if listShowResponses {
- fmt.Println("\nTask Responses:")
- fmt.Println(strings.Repeat("=", 50))
-
- for i, task := range tasks {
- if i >= 5 { // Limit responses to first 5 tasks to avoid overwhelming output
- break
- }
-
- if showingAllCallbacks {
- fmt.Printf("\nTask %d from callback %d (%s):\n", task.DisplayID, task.CallbackID, task.Command)
- } else {
- fmt.Printf("\nTask %d (%s):\n", task.DisplayID, task.Command)
- }
- fmt.Println(strings.Repeat("-", 30))
-
- // For all callbacks mode, task already includes response from GetAllTasksWithResponses
- if showingAllCallbacks && task.Response != "" {
- fmt.Println(task.Response)
- } else if !showingAllCallbacks {
- // For specific callback, get full task with response
- fullTask, err := client.GetTaskResponse(ctx, task.ID)
- if err != nil {
- fmt.Printf("Error getting response: %v\n", err)
- continue
- }
-
- if fullTask.Response != "" {
- fmt.Println(fullTask.Response)
- } else {
- fmt.Println("(No response yet)")
- }
- } else {
- fmt.Println("(No response yet)")
- }
- fmt.Println()
- }
- }
-
return nil
}
\ No newline at end of file
cmd/task_view.go
@@ -3,8 +3,10 @@ package cmd
import (
"context"
"fmt"
+ "mysh/pkg/cache"
"mysh/pkg/mythic"
"mysh/pkg/mythic/api"
+ "os"
"strconv"
"strings"
"time"
@@ -46,15 +48,46 @@ func runTaskView(cmd *cobra.Command, args []string) error {
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
- task, err := client.GetTaskResponse(ctx, taskID)
+ // Initialize cache
+ taskCache, err := cache.New(mythicURL)
if err != nil {
- return fmt.Errorf("failed to get task %d: %w", taskID, err)
+ // If cache initialization fails, continue without caching
+ taskCache = nil
+ }
+
+ var task *mythic.Task
+
+ // Try to get from cache first for completed tasks
+ if taskCache != nil {
+ if cachedTask, found := taskCache.GetCachedTask(taskID, mythicURL); found {
+ task = cachedTask
+ }
+ }
+
+ // If not in cache, fetch from server
+ if task == nil {
+ task, err = client.GetTaskResponse(ctx, taskID)
+ if err != nil {
+ return fmt.Errorf("failed to get task %d: %w", taskID, err)
+ }
+
+ // Cache the result if it's completed and cache is available
+ if taskCache != nil {
+ if err := taskCache.CacheTask(task, mythicURL); err != nil {
+ // Log error but don't fail the command
+ fmt.Fprintf(os.Stderr, "Warning: failed to cache task result: %v\n", err)
+ }
+ }
}
// If raw output requested, only print the response
if taskViewRawOutput {
if task.Response != "" {
fmt.Print(task.Response)
+ // Ensure output ends with a newline
+ if !strings.HasSuffix(task.Response, "\n") {
+ fmt.Print("\n")
+ }
}
return nil
}
pkg/cache/cache.go
@@ -0,0 +1,241 @@
+package cache
+
+import (
+ "crypto/sha256"
+ "encoding/json"
+ "fmt"
+ "mysh/pkg/mythic"
+ "net/url"
+ "os"
+ "path/filepath"
+ "strings"
+ "time"
+)
+
+// TaskCache manages caching of completed task results
+type TaskCache struct {
+ cacheDir string
+}
+
+// extractHostname safely extracts hostname from a server URL for directory naming
+func extractHostname(serverURL string) string {
+ // Parse the URL to extract hostname
+ parsedURL, err := url.Parse(serverURL)
+ if err != nil {
+ // If parsing fails, use a sanitized version of the full URL
+ return sanitizeForPath(serverURL)
+ }
+
+ hostname := parsedURL.Hostname()
+ if hostname == "" {
+ // Fallback to host (includes port if present)
+ hostname = parsedURL.Host
+ }
+
+ if hostname == "" {
+ // Final fallback to sanitized URL
+ return sanitizeForPath(serverURL)
+ }
+
+ return sanitizeForPath(hostname)
+}
+
+// sanitizeForPath removes characters that aren't safe for directory names
+func sanitizeForPath(input string) string {
+ // Replace unsafe characters with underscores
+ unsafe := []string{"/", "\\", ":", "*", "?", "\"", "<", ">", "|", " "}
+ result := input
+ for _, char := range unsafe {
+ result = strings.ReplaceAll(result, char, "_")
+ }
+ return result
+}
+
+// New creates a new TaskCache instance using XDG cache directory with server-specific subdirectory
+func New(serverURL string) (*TaskCache, error) {
+ // Use XDG_CACHE_HOME if set, otherwise use default ~/.cache
+ cacheHome := os.Getenv("XDG_CACHE_HOME")
+ if cacheHome == "" {
+ homeDir, err := os.UserHomeDir()
+ if err != nil {
+ return nil, fmt.Errorf("failed to get user home directory: %w", err)
+ }
+ cacheHome = filepath.Join(homeDir, ".cache")
+ }
+
+ // Extract hostname for server-specific directory
+ hostname := extractHostname(serverURL)
+ cacheDir := filepath.Join(cacheHome, "mysh", hostname)
+
+ // Create cache directory if it doesn't exist
+ if err := os.MkdirAll(cacheDir, 0755); err != nil {
+ return nil, fmt.Errorf("failed to create cache directory: %w", err)
+ }
+
+ return &TaskCache{cacheDir: cacheDir}, nil
+}
+
+// CachedTask represents a cached task with metadata
+type CachedTask struct {
+ Task *mythic.Task `json:"task"`
+ CachedAt time.Time `json:"cached_at"`
+ ServerURL string `json:"server_url"`
+}
+
+
+// generateCacheKey creates a unique cache key for a task
+func (tc *TaskCache) generateCacheKey(taskID int, serverURL string) string {
+ // Create a hash based on task ID and server URL to ensure uniqueness
+ h := sha256.New()
+ h.Write([]byte(fmt.Sprintf("%d:%s", taskID, serverURL)))
+ return fmt.Sprintf("task_%d_%x.json", taskID, h.Sum(nil)[:8])
+}
+
+
+// GetCachedTask retrieves a cached task if it exists and is for a completed task
+func (tc *TaskCache) GetCachedTask(taskID int, serverURL string) (*mythic.Task, bool) {
+ cacheKey := tc.generateCacheKey(taskID, serverURL)
+ cachePath := filepath.Join(tc.cacheDir, cacheKey)
+
+ // Check if cache file exists
+ if _, err := os.Stat(cachePath); os.IsNotExist(err) {
+ return nil, false
+ }
+
+ // Read cache file
+ file, err := os.Open(cachePath)
+ if err != nil {
+ return nil, false
+ }
+ defer file.Close()
+
+ // Decode cached task
+ var cachedTask CachedTask
+ if err := json.NewDecoder(file).Decode(&cachedTask); err != nil {
+ // If cache is corrupted, remove it
+ os.Remove(cachePath)
+ return nil, false
+ }
+
+ // Verify this cache is for the same server
+ if cachedTask.ServerURL != serverURL {
+ return nil, false
+ }
+
+ // Only return cached results for completed tasks
+ if !cachedTask.Task.Completed && cachedTask.Task.Status != "completed" && cachedTask.Task.Status != "error" {
+ return nil, false
+ }
+
+ return cachedTask.Task, true
+}
+
+// CacheTask stores a completed task result in cache
+func (tc *TaskCache) CacheTask(task *mythic.Task, serverURL string) error {
+ // Only cache completed tasks
+ if !task.Completed && task.Status != "completed" && task.Status != "error" {
+ return nil
+ }
+
+ cacheKey := tc.generateCacheKey(task.ID, serverURL)
+ cachePath := filepath.Join(tc.cacheDir, cacheKey)
+
+ // Create cache entry
+ cachedTask := CachedTask{
+ Task: task,
+ CachedAt: time.Now(),
+ ServerURL: serverURL,
+ }
+
+ // Write to temporary file first, then rename (atomic operation)
+ tempPath := cachePath + ".tmp"
+ file, err := os.Create(tempPath)
+ if err != nil {
+ return fmt.Errorf("failed to create cache file: %w", err)
+ }
+ defer file.Close()
+
+ if err := json.NewEncoder(file).Encode(cachedTask); err != nil {
+ os.Remove(tempPath)
+ return fmt.Errorf("failed to encode cache data: %w", err)
+ }
+
+ // Atomic rename
+ if err := os.Rename(tempPath, cachePath); err != nil {
+ os.Remove(tempPath)
+ return fmt.Errorf("failed to finalize cache file: %w", err)
+ }
+
+ return nil
+}
+
+// CleanOldCache removes cache entries older than the specified duration
+func (tc *TaskCache) CleanOldCache(maxAge time.Duration) error {
+ entries, err := os.ReadDir(tc.cacheDir)
+ if err != nil {
+ return fmt.Errorf("failed to read cache directory: %w", err)
+ }
+
+ cutoff := time.Now().Add(-maxAge)
+
+ for _, entry := range entries {
+ if !entry.IsDir() && filepath.Ext(entry.Name()) == ".json" {
+ cachePath := filepath.Join(tc.cacheDir, entry.Name())
+
+ // Check file modification time
+ info, err := entry.Info()
+ if err != nil {
+ continue
+ }
+
+ if info.ModTime().Before(cutoff) {
+ os.Remove(cachePath)
+ }
+ }
+ }
+
+ return nil
+}
+
+// GetCacheInfo returns information about the cache directory
+func (tc *TaskCache) GetCacheInfo() (string, int, int64, error) {
+ entries, err := os.ReadDir(tc.cacheDir)
+ if err != nil {
+ return tc.cacheDir, 0, 0, fmt.Errorf("failed to read cache directory: %w", err)
+ }
+
+ var totalSize int64
+ fileCount := 0
+
+ for _, entry := range entries {
+ if !entry.IsDir() && filepath.Ext(entry.Name()) == ".json" {
+ info, err := entry.Info()
+ if err == nil {
+ totalSize += info.Size()
+ fileCount++
+ }
+ }
+ }
+
+ return tc.cacheDir, fileCount, totalSize, nil
+}
+
+// ClearCache removes all cached task results
+func (tc *TaskCache) ClearCache() error {
+ entries, err := os.ReadDir(tc.cacheDir)
+ if err != nil {
+ return fmt.Errorf("failed to read cache directory: %w", err)
+ }
+
+ for _, entry := range entries {
+ if !entry.IsDir() && filepath.Ext(entry.Name()) == ".json" {
+ cachePath := filepath.Join(tc.cacheDir, entry.Name())
+ if err := os.Remove(cachePath); err != nil {
+ return fmt.Errorf("failed to remove cache file %s: %w", entry.Name(), err)
+ }
+ }
+ }
+
+ return nil
+}
+
pkg/mythic/client.go
@@ -10,6 +10,7 @@ import (
"net/http"
"net/url"
"regexp"
+ "slices"
"strings"
"unicode/utf8"
@@ -322,6 +323,74 @@ func (c *Client) GetTaskResponse(ctx context.Context, taskID int) (*Task, error)
}, nil
}
+// GetTasksWithResponsesByIDs efficiently fetches multiple tasks with their responses in a single query
+func (c *Client) GetTasksWithResponsesByIDs(ctx context.Context, taskIDs []int) ([]Task, error) {
+ if len(taskIDs) == 0 {
+ return []Task{}, nil
+ }
+
+ req := graphql.NewRequest(GetTasksWithResponsesByIDs)
+ req.Var("task_ids", taskIDs)
+ req.Header.Set("Authorization", "Bearer "+c.token)
+
+ var resp struct {
+ Task []struct {
+ ID int `json:"id"`
+ DisplayID int `json:"display_id"`
+ CommandName string `json:"command_name"`
+ OriginalParams string `json:"original_params"`
+ DisplayParams string `json:"display_params"`
+ Status string `json:"status"`
+ Completed bool `json:"completed"`
+ Timestamp string `json:"timestamp"`
+ Callback struct {
+ ID int `json:"id"`
+ DisplayID int `json:"display_id"`
+ Host string `json:"host"`
+ User string `json:"user"`
+ } `json:"callback"`
+ Responses []struct {
+ ResponseText string `json:"response_text"`
+ Timestamp string `json:"timestamp"`
+ } `json:"responses"`
+ } `json:"task"`
+ }
+
+ if err := c.client.Run(ctx, req, &resp); err != nil {
+ return nil, fmt.Errorf("failed to get tasks with responses: %w", err)
+ }
+
+ tasks := make([]Task, 0, len(resp.Task))
+ for _, task := range resp.Task {
+ // Concatenate all response chunks in chronological order
+ var responseBuilder strings.Builder
+ for _, respChunk := range task.Responses {
+ decodedText := decodeResponseText(respChunk.ResponseText)
+ responseBuilder.WriteString(decodedText)
+ }
+ response := responseBuilder.String()
+
+ status := task.Status
+ if task.Completed {
+ status = "completed"
+ }
+
+ tasks = append(tasks, Task{
+ ID: task.ID,
+ DisplayID: task.DisplayID,
+ Command: task.CommandName,
+ Params: task.DisplayParams,
+ Status: status,
+ Response: response,
+ CallbackID: task.Callback.ID,
+ Completed: task.Completed,
+ Timestamp: task.Timestamp,
+ })
+ }
+
+ return tasks, nil
+}
+
func (c *Client) GetCallbackTasks(ctx context.Context, callbackID int) ([]Task, error) {
req := graphql.NewRequest(GetCallbackTasks)
@@ -602,6 +671,62 @@ func (c *Client) GetTasksWithResponses(ctx context.Context, callbackID int, limi
return tasks, nil
}
+func (c *Client) GetAllTasks(ctx context.Context, limit int) ([]Task, error) {
+ req := graphql.NewRequest(GetAllTasks)
+
+ if limit > 0 {
+ req.Var("limit", limit)
+ }
+ req.Header.Set("Authorization", "Bearer "+c.token)
+
+ var resp struct {
+ Task []struct {
+ ID int `json:"id"`
+ DisplayID int `json:"display_id"`
+ CommandName string `json:"command_name"`
+ OriginalParams string `json:"original_params"`
+ DisplayParams string `json:"display_params"`
+ Status string `json:"status"`
+ Completed bool `json:"completed"`
+ Timestamp string `json:"timestamp"`
+ Callback struct {
+ ID int `json:"id"`
+ DisplayID int `json:"display_id"`
+ Host string `json:"host"`
+ User string `json:"user"`
+ } `json:"callback"`
+ } `json:"task"`
+ }
+
+ if err := c.client.Run(ctx, req, &resp); err != nil {
+ return nil, fmt.Errorf("failed to get all tasks: %w", err)
+ }
+
+ var tasks []Task
+ for _, t := range resp.Task {
+ status := t.Status
+ if t.Completed {
+ status = "completed"
+ }
+
+ tasks = append(tasks, Task{
+ ID: t.ID,
+ DisplayID: t.DisplayID,
+ Command: t.CommandName,
+ Params: t.DisplayParams,
+ Status: status,
+ CallbackID: t.Callback.ID,
+ Timestamp: t.Timestamp,
+ Completed: t.Completed,
+ })
+ }
+
+ // Reverse to show newest (highest ID) last
+ slices.Reverse(tasks)
+
+ return tasks, nil
+}
+
func (c *Client) GetAllTasksWithResponses(ctx context.Context, limit int) ([]Task, error) {
req := graphql.NewRequest(GetAllTasksWithResponsesFromAllCallbacks)
pkg/mythic/queries.go
@@ -380,6 +380,30 @@ query GetTask($id: Int!) {
}
}`
+const GetTasksWithResponsesByIDs = `
+query GetTasksWithResponsesByIDs($task_ids: [Int!]!) {
+ task(where: {id: {_in: $task_ids}}) {
+ id
+ display_id
+ command_name
+ original_params
+ display_params
+ status
+ completed
+ timestamp
+ callback {
+ id
+ display_id
+ host
+ user
+ }
+ responses(order_by: {id: asc}) {
+ response_text
+ timestamp
+ }
+ }
+}`
+
const GetTaskResponses = `
query GetTaskResponses($task_display_id: Int!) {
response(where: {task: {display_id: {_eq: $task_display_id}}}, order_by: {id: asc}) {
@@ -550,6 +574,26 @@ query GetTasksWithResponses($callback_id: Int, $limit: Int) {
}
}`
+const GetAllTasks = `
+query GetAllTasks($limit: Int) {
+ task(order_by: {id: desc}, limit: $limit) {
+ id
+ display_id
+ command_name
+ original_params
+ display_params
+ status
+ completed
+ timestamp
+ callback {
+ id
+ display_id
+ host
+ user
+ }
+ }
+}`
+
const GetAllTasksWithResponsesFromAllCallbacks = `
query GetAllTasksWithResponsesFromAllCallbacks($limit: Int) {
task(order_by: {id: desc}, limit: $limit) {
TODO.md
@@ -1,17 +1,20 @@
-- ✅ grep should also look at the command execution strings too, not just output
-- ✅ Renamed project from go-mythic to mysh (pronounced mɪʃh)
-- ✅ task-list --all (same as --limit 0), and -a
-- ✅ task-list without callback lists all tasks and shows callback column
-- ✅ task with --no-wait, doesn't wait for output
-- ✅ task [callback] arg make it accept lists, comma seperated, or ranges e.g. 1,2,3-5,10
-- ✅ make cb and callback aliases for callbacks
-- ✅ update callbacks to show last-checkin as 0s or 0m00s ago
-- ✅ make tv --raw have short -r
-- ✅ make tv have more task details in the header
-- ✅ make tv have a -d --details to only show the header
-- make callbacks sortable with (--rev reverse) --id, --type, --host, --user, --process, --last, --description all exclusive - default should be --id with newest (highest) last
-- create a .cache/mysh/ to store task result cached values for completed tasks - use whatever the right XDG env default dir should be
-- invert cobra-cli nto functions (no global vars)
-- update output table formats to respect terminal width
-- MYTHIC_API_INSECURE= boolean --insecure flag - and invert default = false
- forge payloads - equivalency with execute_assembly and inline_assembly
+- MYTHIC_API_INSECURE= boolean --insecure flag - and invert default = false
+- update output table formats to respect terminal width
+- invert cobra-cli nto functions (no global vars)
+- ✅ use per-server cache dirs
+- ✅ make sure raw output has a newline at the end
+- ✅ make mysh cache default to info output
+- ✅ create a .cache/mysh/<server> to store task results
+- ✅ make callbacks sortable by columns
+- ✅ make tv have a -d --details to only show the header
+- ✅ make tv have more task details in the header
+- ✅ make tv --raw have short -r
+- ✅ update callbacks to show last-checkin as 0s or 0m00s ago
+- ✅ make cb and callback aliases for callbacks
+- ✅ task [callback] arg make it accept lists, comma seperated, or ranges e.g. 1,2,3-5,10
+- ✅ task with --no-wait, doesn't wait for output
+- ✅ task-list without callback lists all tasks and shows callback column
+- ✅ task-list --all (same as --limit 0), and -a
+- ✅ Renamed project from go-mythic to mysh (pronounced mɪʃh)
+- ✅ grep should also look at the command execution strings too, not just output