feat: updated cache implementation and fixed bugs

This commit is contained in:
David Allen 2025-06-20 14:14:38 -06:00
parent 17dbfa2b3b
commit e71f33878e
Signed by: towk
GPG key ID: 0430CDBE22619155
7 changed files with 550 additions and 97 deletions

View file

@ -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] <timestamp>",
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")
}

View file

@ -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)
}

176
internal/cache/edit.go vendored Normal file
View file

@ -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
}

View file

@ -1 +0,0 @@
package cache

View file

@ -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 {

60
internal/cache/state.go vendored Normal file
View file

@ -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)
}

View file

@ -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
}