main
Raw Download raw file
  1package cmd
  2
  3import (
  4	"fmt"
  5	"os"
  6
  7	"github.com/jedib0t/go-pretty/v6/table"
  8	"github.com/spf13/cobra"
  9	"golang.org/x/term"
 10)
 11
 12var (
 13	mythicURL  string
 14	token      string
 15	insecure   bool
 16	socksProxy string
 17)
 18
 19var rootCmd = &cobra.Command{
 20	Use:   "mysh",
 21	Short: "A CLI tool for interacting with Mythic C2 framework",
 22	Long: `mysh (pronounced mɪʃh) is a command-line interface for interacting with the Mythic C2 framework.
 23It provides programmatic access to Mythic operations, replacing the need for UI interactions.`,
 24}
 25
 26func Execute() {
 27	err := rootCmd.Execute()
 28	if err != nil {
 29		os.Exit(1)
 30	}
 31}
 32
 33func init() {
 34	rootCmd.PersistentFlags().StringVar(&mythicURL, "url", "", "Mythic GraphQL API URL")
 35	rootCmd.PersistentFlags().StringVar(&token, "token", "", "JWT authentication token")
 36	rootCmd.PersistentFlags().BoolVar(&insecure, "insecure", true, "Skip TLS certificate verification")
 37	rootCmd.PersistentFlags().StringVar(&socksProxy, "socks", "", "SOCKS5 proxy address (e.g., 127.0.0.1:9050)")
 38
 39	// Set default URL from environment if available
 40	if urlEnv := os.Getenv("MYTHIC_API_URL"); urlEnv != "" {
 41		mythicURL = urlEnv
 42	}
 43
 44	// Set default token from environment if available
 45	if tokenEnv := os.Getenv("MYTHIC_API_TOKEN"); tokenEnv != "" {
 46		token = tokenEnv
 47	}
 48
 49	// Set default SOCKS proxy from environment if available
 50	if socksEnv := os.Getenv("MYTHIC_SOCKS"); socksEnv != "" {
 51		socksProxy = socksEnv
 52	}
 53}
 54
 55func validateConfig() error {
 56	if mythicURL == "" {
 57		return fmt.Errorf("Mythic server URL is required. Use --url flag or set MYTHIC_API_URL environment variable")
 58	}
 59	if token == "" {
 60		return fmt.Errorf("authentication token is required. Use --token flag or set MYTHIC_API_TOKEN environment variable")
 61	}
 62	return nil
 63}
 64
 65// getTerminalWidth returns the current terminal width, with a fallback to 80
 66func getTerminalWidth() int {
 67	if width, _, err := term.GetSize(int(os.Stdout.Fd())); err == nil && width > 0 {
 68		return width
 69	}
 70	return 80 // fallback width
 71}
 72
 73// ColumnInfo represents a table column with its properties
 74type ColumnInfo struct {
 75	Header   string
 76	Number   int
 77	MinWidth int
 78	MaxWidth int
 79	Priority int // 1 = highest priority (never drop), higher numbers can be dropped first
 80}
 81
 82// DynamicTable wraps a table with information about visible columns
 83type DynamicTable struct {
 84	table.Writer
 85	VisibleColumns []ColumnInfo
 86}
 87
 88// configureTableForTerminal sets up a table with dynamic column sizing based on terminal width
 89func configureTableForTerminal(t table.Writer, columns []ColumnInfo) *DynamicTable {
 90	termWidth := getTerminalWidth()
 91
 92	// First pass: estimate how many columns we can fit
 93	// The table library adds 2 spaces around each column's content (1 before, 1 after)
 94	// Plus we need some margin for the table structure itself
 95	estimatedColumns := len(columns)
 96	estimatedPadding := (estimatedColumns * 2) + 2 // 2 per column + 2 for margins
 97	availableWidth := termWidth - estimatedPadding
 98
 99	// Filter columns that can fit and distribute width
100	visibleColumns, columnWidths := distributeColumnWidths(columns, availableWidth)
101
102	// Recalculate with actual column count
103	actualPadding := (len(visibleColumns) * 2) + 2
104	availableWidth = termWidth - actualPadding
105
106	// Re-distribute if needed with correct padding
107	if actualPadding < estimatedPadding {
108		visibleColumns, columnWidths = distributeColumnWidths(columns, availableWidth)
109	}
110
111	// Create headers and column configs for visible columns
112	var headers []interface{}
113	var configs []table.ColumnConfig
114
115	for i, col := range visibleColumns {
116		headers = append(headers, col.Header)
117		configs = append(configs, table.ColumnConfig{
118			Number:           i + 1,           // Use sequential numbers for visible columns
119			WidthMin:         columnWidths[i], // Force minimum width
120			WidthMax:         columnWidths[i], // Force maximum width
121			WidthMaxEnforcer: truncateWithEllipsis,
122		})
123	}
124
125	t.AppendHeader(table.Row(headers))
126	t.SetColumnConfigs(configs)
127	// Don't set AllowedRowLength since we're managing column widths precisely
128
129	return &DynamicTable{
130		Writer:         t,
131		VisibleColumns: visibleColumns,
132	}
133}
134
135// AppendRowForColumns adds a row to the table using only the visible columns
136func (dt *DynamicTable) AppendRowForColumns(allData []interface{}) {
137	var visibleData []interface{}
138	for _, col := range dt.VisibleColumns {
139		if col.Number-1 < len(allData) {
140			visibleData = append(visibleData, allData[col.Number-1])
141		}
142	}
143	dt.AppendRow(table.Row(visibleData))
144}
145
146// distributeColumnWidths filters columns by priority and distributes available width
147func distributeColumnWidths(columns []ColumnInfo, availableWidth int) ([]ColumnInfo, []int) {
148	// Sort by priority (lower number = higher priority)
149	sortedCols := make([]ColumnInfo, len(columns))
150	copy(sortedCols, columns)
151
152	// Simple bubble sort by priority
153	for i := 0; i < len(sortedCols)-1; i++ {
154		for j := 0; j < len(sortedCols)-i-1; j++ {
155			if sortedCols[j].Priority > sortedCols[j+1].Priority {
156				sortedCols[j], sortedCols[j+1] = sortedCols[j+1], sortedCols[j]
157			}
158		}
159	}
160
161	// Step 1: Find which columns can fit using minimum widths
162	// Try to fit columns in priority order (lowest priority number = highest priority)
163	var visibleColumns []ColumnInfo
164	totalMinWidth := 0
165
166	for _, col := range sortedCols {
167		// Check if adding this column would exceed available width
168		newTotal := totalMinWidth + col.MinWidth
169		if len(visibleColumns) > 0 {
170			newTotal++ // Add 1 space between columns
171		}
172		if newTotal <= availableWidth {
173			visibleColumns = append(visibleColumns, col)
174			totalMinWidth = newTotal
175		} else {
176			// Can't fit this column or any remaining lower priority columns
177			break
178		}
179	}
180
181	// Ensure we have at least one column (highest priority)
182	if len(visibleColumns) == 0 && len(sortedCols) > 0 {
183		visibleColumns = append(visibleColumns, sortedCols[0])
184		totalMinWidth = sortedCols[0].MinWidth
185	}
186
187	// Step 2: Distribute available width across visible columns
188	columnWidths := make([]int, len(visibleColumns))
189
190	// Start with minimum widths (totalMinWidth already includes padding)
191	actualContentWidth := 0
192	for i, col := range visibleColumns {
193		columnWidths[i] = col.MinWidth
194		actualContentWidth += col.MinWidth
195	}
196
197	// Calculate remaining width to distribute (subtract actual content width, not totalMinWidth with padding)
198	remainingWidth := availableWidth - totalMinWidth
199
200	if remainingWidth > 0 {
201		// Calculate total expansion capacity for proportional distribution
202		totalExpansion := 0
203		for _, col := range visibleColumns {
204			expansion := col.MaxWidth - col.MinWidth
205			if expansion > 0 {
206				totalExpansion += expansion
207			}
208		}
209
210		if totalExpansion > 0 {
211			// Distribute remaining width proportionally
212			for i, col := range visibleColumns {
213				expansion := col.MaxWidth - col.MinWidth
214				if expansion > 0 {
215					extraWidth := (remainingWidth * expansion) / totalExpansion
216					columnWidths[i] = col.MinWidth + extraWidth
217					// Ensure we don't exceed max width
218					if columnWidths[i] > col.MaxWidth {
219						columnWidths[i] = col.MaxWidth
220					}
221				}
222			}
223
224			// Distribute any leftover width (due to rounding or max width constraints)
225			actualUsed := 0
226			for _, width := range columnWidths {
227				actualUsed += width
228			}
229			leftover := availableWidth - actualUsed
230
231			// Only distribute leftover if it's positive and reasonable
232			if leftover > 0 && leftover < len(visibleColumns)*2 {
233				for i := 0; i < leftover && i < len(visibleColumns); i++ {
234					if columnWidths[i] < visibleColumns[i].MaxWidth {
235						columnWidths[i]++
236					}
237				}
238			}
239		}
240	}
241
242	return visibleColumns, columnWidths
243}
244
245// filterColumnsForWidth removes low-priority columns until the table fits in available width
246// Kept for backwards compatibility, but now uses distributeColumnWidths internally
247func filterColumnsForWidth(columns []ColumnInfo, availableWidth int) []ColumnInfo {
248	visibleColumns, _ := distributeColumnWidths(columns, availableWidth)
249	return visibleColumns
250}
251
252// truncateWithEllipsis is a helper function for table column width enforcement
253// It truncates text that exceeds maxLen and adds "..." at the end
254func truncateWithEllipsis(col string, maxLen int) string {
255	if len(col) <= maxLen {
256		return col
257	}
258	if maxLen <= 3 {
259		return col[:maxLen]
260	}
261	return col[:maxLen-3] + "..."
262}