Merge pull request #5 from davidallendj/cache-cmd

Cache cmd

cmd/cache.go

@@ -0,0 +1,96 @@
+package cmd
+
+import (
+	"fmt"
+	"net/url"
+	"os"
+	"strconv"
+
+	magellan "github.com/OpenCHAMI/magellan/internal"
+	"github.com/OpenCHAMI/magellan/internal/cache/sqlite"
+	"github.com/rs/zerolog/log"
+	"github.com/spf13/cobra"
+)
+
+var (
+	withHosts []string
+	withPorts []int
+)
+
+var cacheCmd = &cobra.Command{
+	Use:   "cache",
+	Short: "Manage found assets in cache.",
+	Run: func(cmd *cobra.Command, args []string) {
+		// show the help for cache and exit
+		if len(args) <= 0 {
+			cmd.Help()
+			os.Exit(0)
+		}
+	},
+}
+
+var cacheRemoveCmd = &cobra.Command{
+	Use:   "remove",
+	Short: "Remove a host from a scanned cache list.",
+	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)
+		}
+
+		// 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...)
+	},
+}
+
+func init() {
+	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")
+	cacheCmd.AddCommand(cacheRemoveCmd)
+	rootCmd.AddCommand(cacheCmd)
+}

internal/cache/sqlite/sqlite.go

@@ -2,6 +2,7 @@ package sqlite
 
 import (
 	"fmt"
+	"strings"
 
 	magellan "github.com/OpenCHAMI/magellan/internal"
 	"github.com/OpenCHAMI/magellan/internal/util"
@@ -59,20 +60,40 @@ func InsertScannedAssets(path string, assets ...magellan.RemoteAsset) error {
 	return nil
 }
 
-func DeleteScannedAssets(path string, results ...magellan.RemoteAsset) error {
-	if results == nil {
+func DeleteScannedAssets(path string, assets ...magellan.RemoteAsset) error {
+	var (
+		db  *sqlx.DB
+		tx  *sqlx.Tx
+		err error
+	)
+	if assets == nil {
 		return fmt.Errorf("no assets found")
 	}
-	db, err := sqlx.Open("sqlite3", path)
+	db, err = sqlx.Open("sqlite3", path)
 	if err != nil {
 		return fmt.Errorf("failed to open database: %v", err)
 	}
-	tx := db.MustBegin()
-	for _, state := range results {
-		sql := fmt.Sprintf(`DELETE FROM %s WHERE host = :host, port = :port;`, TABLE_NAME)
-		_, err := tx.NamedExec(sql, &state)
+	tx = db.MustBegin()
+	for _, asset := range assets {
+		// skip if neither host nor port are specified
+		if asset.Host == "" && asset.Port <= 0 {
+			continue
+		}
+		sql := fmt.Sprintf(`DELETE FROM %s`, TABLE_NAME)
+		where := []string{}
+		if asset.Port > 0 {
+			where = append(where, "port=:port")
+		}
+		if asset.Host != "" {
+			where = append(where, "host=:host")
+		}
+		if len(where) <= 0 {
+			continue
+		}
+		sql += fmt.Sprintf(" WHERE %s;", strings.Join(where, " AND "))
+		_, err := tx.NamedExec(sql, &asset)
 		if err != nil {
-			fmt.Printf("failed to execute transaction: %v\n", err)
+			fmt.Printf("failed to execute DELETE transaction: %v\n", err)
 		}
 	}
 

internal/scan.go

@@ -16,11 +16,11 @@ import (
 )
 
 type RemoteAsset struct {
-	Host      string    `json:"host"`
-	Port      int       `json:"port"`
-	Protocol  string    `json:"protocol"`
-	State     bool      `json:"state"`
-	Timestamp time.Time `json:"timestamp"`
+	Host      string    `db:"host" json:"host"`
+	Port      int       `db:"port" json:"port"`
+	Protocol  string    `db:"protocol" json:"protocol"`
+	State     bool      `db:"state" json:"state"`
+	Timestamp time.Time `db:"timestamp" json:"timestamp"`
 }
 
 // ScanParams is a collection of commom parameters passed to the CLI