From a9d59ee50d96c7986bd0589dd931453c97960a16 Mon Sep 17 00:00:00 2001 From: David Allen Date: Fri, 20 Jun 2025 14:14:38 -0600 Subject: [PATCH] feat: updated cache implementation and fixed bugs --- cmd/cache.go | 177 ++++++++++++++++++------------ internal/cache/cache.go | 1 + internal/cache/edit.go | 176 ++++++++++++++++++++++++++++++ internal/cache/modify.go | 1 - internal/cache/sqlite/sqlite.go | 46 +++++++- internal/cache/state.go | 60 +++++++++++ internal/cache/table.go | 186 ++++++++++++++++++++++++++++---- 7 files changed, 550 insertions(+), 97 deletions(-) create mode 100644 internal/cache/edit.go delete mode 100644 internal/cache/modify.go create mode 100644 internal/cache/state.go diff --git a/cmd/cache.go b/cmd/cache.go index a6f7162..ef04376 100644 --- a/cmd/cache.go +++ b/cmd/cache.go @@ -2,10 +2,12 @@ package cmd import ( "fmt" - "net/url" "os" "strconv" + "strings" + "time" + "github.com/araddon/dateparse" "github.com/charmbracelet/bubbles/table" tea "github.com/charmbracelet/bubbletea" "github.com/charmbracelet/lipgloss" @@ -18,6 +20,7 @@ import ( ) var ( + timestamp time.Time cacheOutputFormat string interactive bool withHosts []string @@ -26,74 +29,103 @@ var ( var cacheCmd = &cobra.Command{ Use: "cache", - Short: "Manage found assets in cache.", + Short: "Manage assets (new, add, remove, edit, etc.) in cache.", +} + +var cacheNewCmd = &cobra.Command{ + Use: "new", + Example: ` // create a new cache database + magellan cache new ./path/to/new_cache.db + `, + Short: "Create a new cache at specified path.", + Run: func(cmd *cobra.Command, args []string) { + + }, +} + +var cacheAddCmd = &cobra.Command{ + Use: "add [host] [port] [protocol] ", + Example: ` // create a new entry in cache + magellan cache add https://172.16.0.105 443,5000 tcp "06/25/2025" + + // create new entry with no timestamp specified with specific cache + magellan cache add https://172.16.0.111 443 tcp --cache ./assets.db + `, + Short: "Add an entry to the specified cache.", + Args: cobra.RangeArgs(3, 4), + Run: func(cmd *cobra.Command, args []string) { + // set positional args to vars + host = args[0] + protocol = args[2] + + // convert ports + for _, portStr := range strings.Split(args[1], ",") { + port, err := strconv.ParseInt(portStr, 10, 16) + if err != nil { + log.Error().Err(err).Msg("failed to convert port value") + continue + } + ports = append(ports, int(port)) + } + + // parse timestamp + if len(args) == 4 { + var err error + timestamp, err = dateparse.ParseAny(args[3]) + if err != nil { + log.Error().Err(err).Msg("failed to parse time arg") + os.Exit(1) + } + } else { + log.Warn().Msg("time not specified...using current time") + timestamp = time.Now() + } + + for _, port := range ports { + // create a new entry in cache using flag arguments + err := sqlite.InsertRemoteAssets(cachePath, magellan.RemoteAsset{ + Host: host, + Port: port, + Protocol: protocol, + Timestamp: timestamp, + }) + if err != nil { + log.Error().Err(err). + Str("path", cachePath). + Str("host", host). + Msg("failed to insert asset into cache") + } + } + + }, } var cacheRemoveCmd = &cobra.Command{ - Use: "remove", - Short: "Remove a host from a scanned cache list.", + Use: "remove [hosts...]", + Example: ` // remove an entry in cache + magellan cache remove https://172.16.0.113 + `, + Args: cobra.MinimumNArgs(1), + Short: "Remove hosts from cache.", Run: func(cmd *cobra.Command, args []string) { - assets := []magellan.RemoteAsset{} - - // add all assets directly from positional args - for _, arg := range args { - var ( - port int - uri *url.URL - err error - ) - uri, err = url.ParseRequestURI(arg) - if err != nil { - log.Error().Err(err).Msg("failed to parse arg") - } - - // convert port to its "proper" type - if uri.Port() == "" { - uri.Host += ":443" - } - port, err = strconv.Atoi(uri.Port()) - if err != nil { - log.Error().Err(err).Msg("failed to convert port to integer type") - } - asset := magellan.RemoteAsset{ - Host: fmt.Sprintf("%s://%s", uri.Scheme, uri.Hostname()), - Port: port, - } - assets = append(assets, asset) + err := sqlite.DeleteRemoteAssetsByHost(cachePath, args...) + if err != nil { + log.Error().Err(err).Str("path", cachePath).Msg("failed to delete assets in cache") } - // Add all assets with specified hosts (same host different different ports) - // This should produce the following SQL: - // DELETE FROM magellan_scanned_assets WHERE host=:host - for _, host := range withHosts { - assets = append(assets, magellan.RemoteAsset{ - Host: host, - Port: -1, - }) - } - - // Add all assets with specified ports (same port different hosts) - // This should produce the following SQL: - // DELETE FROM magellan_scanned_assets WHERE port=:port - for _, port := range withPorts { - assets = append(assets, magellan.RemoteAsset{ - Host: "", - Port: port, - }) - } - if len(assets) <= 0 { - log.Error().Msg("nothing to do") - os.Exit(1) - } - sqlite.DeleteScannedAssets(cachePath, assets...) }, } var cacheEditCmd = &cobra.Command{ Use: "edit", - Example: ` magellan cache edit - magellan cache edit --host https://172.16.0.101 --port 443 --protocol udp - magellan cache edit --host https://172.16.0.101 + Example: ` // star the cache editor + magellan cache edit -i + + // edit a single entry only changing values specified (e.g. port and protocol) + magellan cache edit https://172.16.0.101 --port 443 --protocol udp + + // edit two entries' time stamps + magellan cache edit https://172.16.0.101 https://172.16.0.102 --timestamp 06/25/2025 `, Args: cobra.ExactArgs(0), Short: "Edit existing cache data.", @@ -106,17 +138,17 @@ var cacheEditCmd = &cobra.Command{ if interactive { // load the assets found from scan - scannedResults, err := sqlite.GetScannedAssets(cachePath) + scannedResults, err := sqlite.GetRemoteAssets(cachePath) if err != nil { log.Error().Err(err).Str("path", cachePath).Msg("failed to get scanned assets from cache") } // set columns to cache headers columns = []table.Column{ - {Title: "hosts", Width: 20}, - {Title: "ports", Width: 5}, - {Title: "protocol", Width: 8}, - {Title: "timestamp", Width: 12}, + {Title: "host", Width: 30}, + {Title: "ports", Width: 8}, + {Title: "protocol", Width: 10}, + {Title: "timestamp", Width: 20}, } // set rows to cache data @@ -129,7 +161,7 @@ var cacheEditCmd = &cobra.Command{ }) } - // new table + // create a new table assetsTable := table.New( table.WithColumns(columns), table.WithRows(rows), @@ -150,11 +182,16 @@ var cacheEditCmd = &cobra.Command{ Bold(false) assetsTable.SetStyles(styles) - m := cache.Model{Table: assetsTable} + m := cache.NewModel(cachePath, &assetsTable) if _, err := tea.NewProgram(m, tea.WithAltScreen()).Run(); err != nil { fmt.Println("Error running program:", err) os.Exit(1) } + } else { + // non-interactive editting + // for _, host := range args { + + // } } }, } @@ -169,7 +206,7 @@ var cacheInfoCmd = &cobra.Command{ } func init() { - // remove + // remove row from cache cacheRemoveCmd.Flags().StringSliceVar(&withHosts, "with-hosts", []string{}, "Remove all assets with specified hosts") cacheRemoveCmd.Flags().IntSliceVar(&withPorts, "with-ports", []int{}, "Remove all assets with specified ports") @@ -181,16 +218,18 @@ func init() { cacheInfoCmd.Flags().StringVarP(&cacheOutputFormat, "format", "F", FORMAT_LIST, "Set the output format (list|json|yaml)") cacheCmd.AddCommand( + cacheNewCmd, + cacheAddCmd, cacheRemoveCmd, - cacheInfoCmd, cacheEditCmd, - ListCmd, + cacheInfoCmd, + ListCmd, // magellan cache list (alias for 'magellan list') ) rootCmd.AddCommand(cacheCmd) } func printCacheInfo(format string) { - assets, err := sqlite.GetScannedAssets(cachePath) + assets, err := sqlite.GetRemoteAssets(cachePath) if err != nil { log.Error().Err(err).Str("path", cachePath).Msg("failed to get assets to print cache info") } diff --git a/internal/cache/cache.go b/internal/cache/cache.go index 7b4eea1..b62a5f7 100644 --- a/internal/cache/cache.go +++ b/internal/cache/cache.go @@ -8,6 +8,7 @@ import ( type Cache[T any] interface { CreateIfNotExists(path string) (driver.Connector, error) Insert(path string, data ...T) error + Update(path string, data ...T) error Delete(path string, data ...T) error Get(path string) ([]T, error) } diff --git a/internal/cache/edit.go b/internal/cache/edit.go new file mode 100644 index 0000000..bebfe73 --- /dev/null +++ b/internal/cache/edit.go @@ -0,0 +1,176 @@ +package cache + +import ( + "fmt" + "os" + "slices" + "strconv" + "strings" + "time" + + "github.com/araddon/dateparse" + "github.com/charmbracelet/bubbles/table" + tea "github.com/charmbracelet/bubbletea" + "github.com/cznic/mathutil" + "github.com/davidallendj/magellan/internal/cache/sqlite" + magellan "github.com/davidallendj/magellan/pkg" + "github.com/rs/zerolog/log" +) + +func (m *Model) editSelectedRow(value bool) tea.Cmd { + if value { + row := m.Table.SelectedRow() + for i := range m.inputs { + m.inputs[i].SetValue(row[i]) + } + m.Editor.StartEditting() + } else { + m.Editor.StopEditting() + } + return nil +} + +func (m *Model) updateRow() tea.Cmd { + // get updated values from inputs + updated := make(table.Row, len(m.inputs)) + for i, input := range m.inputs { + updated[i] = input.Value() + } + + // update table for selected row + rows := m.Table.Rows() + rows[m.Table.Cursor()] = updated + m.Table.SetRows(rows) + + // go back to selecting view + m.Editor.StopEditting() + m.displayMessage(fmt.Sprintf("updated row at index %d", m.Table.Cursor()), 3) + return nil +} + +func (m *Model) addRowAfterCursor() tea.Cmd { + position := mathutil.Clamp(m.Table.Cursor()+1, 0, len(m.Table.Rows())-1) + rows := slices.Insert(m.Table.Rows(), position, table.Row{ + "https://127.0.0.1", + "443", + "tcp", + fmt.Sprintf("%d", time.Now().Unix()), + }) + + m.Table.SetRows(rows) + m.displayMessage(fmt.Sprintf("add new row at index %d", m.Table.Cursor()+1), 3) + return nil +} + +func (m *Model) addRowAtEnd() tea.Cmd { + rows := append(m.Table.Rows(), table.Row{ + "https://127.0.0.1", + "443", + "tcp", + fmt.Sprintf("%d", time.Now().Unix()), + }) + m.Table.SetRows(rows) + m.displayMessage(fmt.Sprintf("add new row at index %d", len(m.Table.Rows())), 3) + return nil +} + +func (m *Model) deleteSelectedRow() tea.Cmd { + if len(m.Table.Rows()) == 0 { + m.displayMessage("nothing to delete...", 3) + return nil + } + m.Table.SetRows(slices.Delete(m.Table.Rows(), m.Table.Cursor(), m.Table.Cursor()+1)) + m.displayMessage(fmt.Sprintf("deleted row at index %d", m.Table.Cursor()), 3) + return nil +} + +func (m *Model) updateInputs(msg tea.Msg) tea.Cmd { + cmds := make([]tea.Cmd, len(m.inputs)) + + // Only text inputs with Focus() set will respond, so it's safe to simply + // update all of them here without any further logic. + for i := range m.inputs { + m.inputs[i], cmds[i] = m.inputs[i].Update(msg) + } + + return tea.Batch(cmds...) +} + +func (m *Model) editRowView() string { + display := fmt.Sprintf("Editting row %d...\n\n", m.Table.Cursor()) + + // update the values of each input + const lc = 9 + for i := range m.inputs { + diff := lc - len(m.Table.Columns()[i].Title) + spacing := strings.Repeat(" ", diff) + // m.inputs[i].SetValue(row[i]) + display += fmt.Sprintf("%d: ", i) + display += m.Table.Columns()[i].Title + ": " + display += spacing + m.inputs[i].View() + "\n" + } + + // submit button + // button := &blurredButton + // if m.focusIndex == len(m.inputs) { + // button = &focusedButton + // } + // display += fmt.Sprintf("\n\n%s\n\n", *button) + + // add helper info in footer + display += footerStyle.Render(` + ⬇/⬆: move cursor; + w, enter: save edit; esc: cancel; ctrl+c: quit w/o saving; +`) + + return display +} + +func (m *Model) updateCacheData() tea.Cmd { + // create assets from table data + assets := []magellan.RemoteAsset{} + for _, row := range m.Table.Rows() { + // convert port + port, err := strconv.ParseInt(row[1], 10, 32) + if err != nil { + log.Error().Err(err).Str("host", row[0]).Msg("failed to parse port value") + return nil + } + + // parse timestamp + timestamp, err := dateparse.ParseAny(row[3]) + if err != nil { + log.Error().Err(err).Str("host", row[0]).Msg("failed to parse date/time") + return nil + } + + assets = append(assets, magellan.RemoteAsset{ + Host: row[0], + Port: int(port), + Protocol: row[2], + Timestamp: timestamp, + }) + } + + // remove current database file + err := os.Remove(m.CachePath) + if err != nil { + log.Error().Err(err).Str("path", m.CachePath).Msg("failed to remove old cache") + return nil + } + + // create the file again... + _, err = sqlite.CreateRemoteAssetsIfNotExists(m.CachePath) + if err != nil { + log.Error().Err(err).Str("path", m.CachePath).Msg("failed to create new cache") + return nil + } + + // write assets to database + err = sqlite.InsertRemoteAssets(m.CachePath, assets...) + if err != nil { + log.Error().Err(err).Str("path", m.CachePath).Msg("failed to insert data into cache") + return nil + } + return nil +} diff --git a/internal/cache/modify.go b/internal/cache/modify.go deleted file mode 100644 index 08bf029..0000000 --- a/internal/cache/modify.go +++ /dev/null @@ -1 +0,0 @@ -package cache diff --git a/internal/cache/sqlite/sqlite.go b/internal/cache/sqlite/sqlite.go index bbbdff3..b2966a6 100644 --- a/internal/cache/sqlite/sqlite.go +++ b/internal/cache/sqlite/sqlite.go @@ -11,7 +11,17 @@ import ( const TABLE_NAME = "magellan_scanned_assets" -func CreateScannedAssetIfNotExists(path string) (*sqlx.DB, error) { +func GetColumns() []string { + return []string{ + "host", + "port", + "protocol", + "state", + "timestamp", + } +} + +func CreateRemoteAssetsIfNotExists(path string) (*sqlx.DB, error) { schema := fmt.Sprintf(` CREATE TABLE IF NOT EXISTS %s ( host TEXT NOT NULL, @@ -34,13 +44,13 @@ func CreateScannedAssetIfNotExists(path string) (*sqlx.DB, error) { return db, nil } -func InsertScannedAssets(path string, assets ...magellan.RemoteAsset) error { +func InsertRemoteAssets(path string, assets ...magellan.RemoteAsset) error { if assets == nil { return fmt.Errorf("states == nil") } // create database if it doesn't already exist - db, err := CreateScannedAssetIfNotExists(path) + db, err := CreateRemoteAssetsIfNotExists(path) if err != nil { return err } @@ -62,7 +72,33 @@ func InsertScannedAssets(path string, assets ...magellan.RemoteAsset) error { return nil } -func DeleteScannedAssets(path string, assets ...magellan.RemoteAsset) error { +func DeleteRemoteAssetsByHost(path string, hosts ...string) error { + var ( + db *sqlx.DB + tx *sqlx.Tx + err error + ) + db, err = sqlx.Open("sqlite3", path) + if err != nil { + return fmt.Errorf("failed to open database: %v", err) + } + tx = db.MustBegin() + for _, host := range hosts { + sql := fmt.Sprintf(`DELETE FROM %s WHERE host='%s'`, TABLE_NAME, host) + _, err := tx.Exec(sql, &host) + if err != nil { + fmt.Printf("failed to execute DELETE transaction: %v\n", err) + } + } + + err = tx.Commit() + if err != nil { + return fmt.Errorf("failed to commit transaction: %v", err) + } + return nil +} + +func DeleteRemoteAssets(path string, assets ...magellan.RemoteAsset) error { var ( db *sqlx.DB tx *sqlx.Tx @@ -91,7 +127,7 @@ func DeleteScannedAssets(path string, assets ...magellan.RemoteAsset) error { return nil } -func GetScannedAssets(path string) ([]magellan.RemoteAsset, error) { +func GetRemoteAssets(path string) ([]magellan.RemoteAsset, error) { // check if path exists first to prevent creating the database _, exists := util.PathExists(path) if !exists { diff --git a/internal/cache/state.go b/internal/cache/state.go new file mode 100644 index 0000000..447c88f --- /dev/null +++ b/internal/cache/state.go @@ -0,0 +1,60 @@ +package cache + +type editorState interface{} +type selecting struct{} +type editting struct{} + +type Editor struct { + state editorState +} + +// transition between state + +func NewEditor() Editor { + return Editor{state: selecting{}} +} + +func (e *Editor) StartEditting() editting { + e.state = editting{} + return e.state.(editting) +} + +func (e *Editor) StopEditting() selecting { + e.state = selecting{} + return e.state.(selecting) +} + +func (s editting) FinishEditting(e *Editor) selecting { + e.state = selecting{} + return e.state.(selecting) +} + +func (e *Editor) GetState() editorState { + return e.state +} + +func (e *Editor) GetStateString() string { + switch e.state.(type) { + case selecting: + return "selecting" + case editting: + return "editting" + } + return "" +} + +func (e *Editor) IsEditting() bool { + return e.GetStateString() == "editting" +} + +func (e *Editor) IsSelecting() bool { + return e.GetStateString() == "selecting" +} + +func test() { + editor := NewEditor() + editting := editor.StartEditting() + + // ...later on... + editting.FinishEditting(&editor) +} diff --git a/internal/cache/table.go b/internal/cache/table.go index 777308a..0dfdeea 100644 --- a/internal/cache/table.go +++ b/internal/cache/table.go @@ -1,22 +1,77 @@ package cache import ( + "fmt" + "time" + "github.com/charmbracelet/bubbles/table" + "github.com/charmbracelet/bubbles/textinput" tea "github.com/charmbracelet/bubbletea" "github.com/charmbracelet/lipgloss" + "github.com/cznic/mathutil" ) -var baseStyle = lipgloss.NewStyle(). - BorderStyle(lipgloss.NormalBorder()). - BorderForeground(lipgloss.Color("240")) +var ( + baseStyle = lipgloss.NewStyle(). + BorderStyle(lipgloss.NormalBorder()). + BorderForeground(lipgloss.Color("240")) + footerStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("241")) + focusedStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("205")) + blurredStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + cursorStyle = focusedStyle + noStyle = lipgloss.NewStyle() + helpStyle = blurredStyle + cursorModeHelpStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("244")) + + focusedButton = focusedStyle.Render("[ Update ]") + blurredButton = fmt.Sprintf("[ %s ]", blurredStyle.Render("Update")) +) type Model struct { - selected int - Table table.Model + CachePath string + Editor Editor + Table table.Model + + // privates + inputs []textinput.Model + focusIndex int + statusMsg string + statusTimer time.Duration } -func NewModel() Model { - return Model{} +func NewModel(cachePath string, table *table.Model) Model { + m := Model{ + CachePath: cachePath, + Editor: NewEditor(), + Table: *table, + inputs: make([]textinput.Model, 4), + focusIndex: 0, + } + for i := range len(m.inputs) { + input := textinput.New() + input.Cursor.Style = cursorStyle + input.CharLimit = 256 + switch i { + case 0: + input.Placeholder = "host" + input.Focus() + input.PromptStyle = focusedStyle + input.TextStyle = focusedStyle + case 1: + input.Placeholder = "ports" + input.CharLimit = 5 + case 2: + input.Placeholder = "protocol" + input.CharLimit = 5 + + case 3: + input.Placeholder = "timestamp" + input.CharLimit = 10 + } + m.inputs[i] = input + } + + return m } func (m Model) Init() tea.Cmd { return nil } @@ -25,28 +80,115 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { var cmd tea.Cmd switch msg := msg.(type) { case tea.WindowSizeMsg: - // m.Table = m.Table.Width(msg.Width) - // m.Table = m.Table.Height(msg.Height) + // m.Table.SetWidth(msg.Width) + // m.Table.SetHeight(msg.Height) case tea.KeyMsg: + // do always regardless of state switch msg.String() { - case "esc": - if m.Table.Focused() { - m.Table.Blur() - } else { - m.Table.Focus() - } - case "q", "ctrl+c": + case "ctrl+c": return m, tea.Quit - case "enter": - return m, tea.Batch( - tea.Printf("Selected host '%s'", m.Table.SelectedRow()[0]), - ) } + + // do while editting table row + if m.Editor.IsEditting() { + switch msg.String() { + case "esc": + return m, m.editSelectedRow(false) + case "w", "enter": + return m, m.updateRow() + // case "enter": + // // Did the user press enter while the submit button was focused? + // // If so, exit. + // if m.focusIndex == len(m.inputs) { + // return m, tea.Quit + // } + + case "tab", "shift+tab", "up", "down": + s := msg.String() + switch s { + case "up", "shift+tab": + m.focusIndex = mathutil.Clamp(m.focusIndex-1, 0, len(m.inputs)) + case "down", "tab": + m.focusIndex = mathutil.Clamp(m.focusIndex+1, 0, len(m.inputs)) + } + + cmds := make([]tea.Cmd, len(m.inputs)) + for i := range m.inputs { + if i == m.focusIndex { + // Set focused state + cmds[i] = m.inputs[i].Focus() + m.inputs[i].PromptStyle = focusedStyle + m.inputs[i].TextStyle = focusedStyle + continue + } + // Remove focused state + m.inputs[i].Blur() + m.inputs[i].PromptStyle = noStyle + m.inputs[i].TextStyle = noStyle + } + return m, tea.Batch(cmds...) + } + cmd = m.updateInputs(msg) + return m, tea.Batch(cmd) + // do while making selection from table + } else if m.Editor.IsSelecting() { + switch msg.String() { + case "enter": + return m, m.editSelectedRow(true) + case "esc": + return m, tea.Quit + case "ctrl+w": + m.updateCacheData() + return m, tea.Quit + case "w", "ctrl+s": + return m, m.updateCacheData() + case "d": + return m, m.deleteSelectedRow() + case "a": + return m, m.addRowAfterCursor() + case "ctrl+a": + return m, m.addRowAtEnd() + case "q": + return m, tea.Quit + } + m.Table, cmd = m.Table.Update(msg) + } + } - m.Table, cmd = m.Table.Update(msg) return m, cmd } func (m Model) View() string { - return baseStyle.Render(m.Table.View()) + "\n" + if m.Editor.IsEditting() { + return m.editRowView() + } else if m.Editor.IsSelecting() { + display := m.Table.View() + "\n" + display += footerStyle.Render(fmt.Sprintf(` + j/k, ⬇/⬆ : move cursor; enter: choose; + a: add row after cursor; ctrl+a: add row at end + d: delete row at cursor; + w, ctrl+s: save data; + ctrl+w: save and quit; q, esc: quit w/o saving; +---------------------------------------------------------- + cache: %s + cursor: %d, rows: %d + %s +`, m.CachePath, m.Table.Cursor(), len(m.Table.Rows()), m.statusMsg)) + + return baseStyle.Render(display + "\n") + } + return "Something went wrong..." +} + +func (m *Model) displayMessage(message string, timeout time.Duration) { + + // show the message arg for alloted time + m.statusMsg = message + reset := make(chan string) + go func() { + time.Sleep(timeout) + reset <- "" + }() + m.statusMsg = <-reset + }