From 0d6cfdec2b53aa730b051e39543fb5273db3785e Mon Sep 17 00:00:00 2001 From: David Allen Date: Fri, 20 Jun 2025 15:07:02 -0600 Subject: [PATCH] feat: add non-interactive cache editting --- cmd/cache.go | 87 ++++++++++++++++++++++++++++----- cmd/crawl.go | 2 +- internal/cache/sqlite/sqlite.go | 21 ++++++++ 3 files changed, 97 insertions(+), 13 deletions(-) diff --git a/cmd/cache.go b/cmd/cache.go index ef04376..316be8d 100644 --- a/cmd/cache.go +++ b/cmd/cache.go @@ -20,11 +20,10 @@ import ( ) var ( + timestampf string timestamp time.Time cacheOutputFormat string interactive bool - withHosts []string - withPorts []int ) var cacheCmd = &cobra.Command{ @@ -127,7 +126,20 @@ var cacheEditCmd = &cobra.Command{ // 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), + Args: func(cmd *cobra.Command, args []string) error { + if interactive { + // must have no args if interactive + if err := cobra.ExactArgs(0)(cmd, args); err != nil { + return err + } + } else { + // must have at least one arg if not interactive + if err := cobra.MinimumNArgs(1)(cmd, args); err != nil { + return err + } + } + return nil + }, Short: "Edit existing cache data.", Run: func(cmd *cobra.Command, args []string) { var ( @@ -189,9 +201,62 @@ var cacheEditCmd = &cobra.Command{ } } else { // non-interactive editting - // for _, host := range args { + for _, host := range args { + // get the asset from cache for host + asset, err := sqlite.GetRemoteAsset(cachePath, host) + if err != nil { + log.Warn().Err(err). + Str("host", host). + Str("path", cachePath). + Msg("failed to get asset from cache") + continue + } + if asset == nil { + log.Warn().Err(err). + Str("host", host). + Str("path", cachePath). + Msg("found asset is not valid") + continue + } - // } + // only modify values that are set + if host != "" { + asset.Host = host + } + if protocol != "" { + asset.Protocol = protocol + } + if timestampf != "" { + newTimestamp, err := dateparse.ParseAny(timestampf) + if err != nil { + log.Error().Err(err).Msg("failed to parse timestamp value") + } else { + asset.Timestamp = newTimestamp + } + } + + // reinsert the asset into cache for each port + for _, port := range ports { + newAsset := *asset + newAsset.Port = port + err = sqlite.DeleteRemoteAssetsByHost(cachePath, host) + if err != nil { + log.Error().Err(err). + Str("host", host). + Str("path", cachePath). + Msg("failed to delete asset in cache") + continue + } + err = sqlite.InsertRemoteAssets(cachePath, newAsset) + if err != nil { + log.Error().Err(err). + Str("host", host). + Str("path", cachePath). + Msg("failed to re-insert asset into cache") + } + } + + } } }, } @@ -206,13 +271,11 @@ var cacheInfoCmd = &cobra.Command{ } func init() { - // 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") - - cacheEditCmd.Flags().IntSliceVar(&ports, "port", nil, "Adds additional ports to scan for each host with unspecified ports.") - cacheEditCmd.Flags().StringVar(&scheme, "scheme", "https", "Set the default scheme to use if not specified in host URI. (default is 'https')") - cacheEditCmd.Flags().StringVar(&protocol, "protocol", "tcp", "Set the default protocol to use in scan. (default is 'tcp')") + cacheEditCmd.Flags().StringVar(&host, "host", "", "Set the new host value.") + cacheEditCmd.Flags().IntSliceVar(&ports, "port", nil, "Set the new port values as comma-separated list.") + cacheEditCmd.Flags().StringVar(&scheme, "scheme", "https", "Set the new scheme value. (default is 'https')") + cacheEditCmd.Flags().StringVar(&protocol, "protocol", "tcp", "Set the new protocol value. (default is 'tcp')") + cacheEditCmd.Flags().StringVar(×tampf, "timestamp", "", "Set the new timestamp value.") cacheEditCmd.Flags().BoolVarP(&interactive, "interactive", "i", false, "Start an interactive TUI to edit cache data") cacheInfoCmd.Flags().StringVarP(&cacheOutputFormat, "format", "F", FORMAT_LIST, "Set the output format (list|json|yaml)") diff --git a/cmd/crawl.go b/cmd/crawl.go index 038e6be..4929978 100644 --- a/cmd/crawl.go +++ b/cmd/crawl.go @@ -30,7 +30,7 @@ var CrawlCmd = &cobra.Command{ Args: func(cmd *cobra.Command, args []string) error { // Validate that the only argument is a valid URI var err error - if err := cobra.ExactArgs(1)(cmd, args); err != nil { + if err = cobra.ExactArgs(1)(cmd, args); err != nil { return err } args[0], err = urlx.Sanitize(args[0]) diff --git a/internal/cache/sqlite/sqlite.go b/internal/cache/sqlite/sqlite.go index b2966a6..b7be903 100644 --- a/internal/cache/sqlite/sqlite.go +++ b/internal/cache/sqlite/sqlite.go @@ -147,3 +147,24 @@ func GetRemoteAssets(path string) ([]magellan.RemoteAsset, error) { } return results, nil } + +func GetRemoteAsset(path string, host string) (*magellan.RemoteAsset, error) { + // check if path exists first to prevent creating the database + _, exists := util.PathExists(path) + if !exists { + return nil, fmt.Errorf("no file found") + } + + // now check if the file is the SQLite database + db, err := sqlx.Open("sqlite3", path) + if err != nil { + return nil, fmt.Errorf("failed to open database: %v", err) + } + + results := []magellan.RemoteAsset{} + err = db.Select(&results, fmt.Sprintf("SELECT * FROM %s ORDER BY host ASC, port ASC;", TABLE_NAME)) + if err != nil { + return nil, fmt.Errorf("failed to retrieve assets: %v", err) + } + return &results[0], nil +}