updated update check

Co-Authored-By: Naomi Kramer <[email protected]>
activecm · Jun 28, 2024 · 268ab78 · 268ab78
1 parent 86c60ce
commit 268ab78
Show file tree

Hide file tree

Showing 7 changed files with 67 additions and 18 deletions.
diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -1,8 +1,12 @@
 package cmd
 
 import (
+	"activecm/rita/config"
+	"activecm/rita/util"
 	"errors"
+	"fmt"
 
+	"github.com/google/go-github/github"
 	"github.com/spf13/afero"
 	"github.com/urfave/cli/v2"
 )
@@ -33,3 +37,26 @@ func ConfigFlag(required bool) *cli.StringFlag {
 		},
 	}
 }
+
+func CheckForUpdate(cCtx *cli.Context, afs afero.Fs) error {
+	// get the current version
+	currentVersion := config.Version
+
+	// load config file
+	cfg, err := config.LoadConfig(afs, cCtx.String("config"))
+	if err != nil {
+		return fmt.Errorf("error loading config file: %w", err)
+	}
+
+	// check for update if version is set
+	if cfg.UpdateCheckEnabled && currentVersion != "" {
+		newer, latestVersion, err := util.CheckForNewerVersion(github.NewClient(nil), "v0.0.0")
+		if err != nil {
+			return fmt.Errorf("error checking for newer version of RITA: %w", err)
+		}
+		if newer {
+			fmt.Printf("\n\t✨ A newer version (%s) of RITA is available! https://github.com/activecm/rita/releases ✨\n\n", latestVersion)
+		}
+	}
+	return nil
+}
diff --git a/cmd/delete.go b/cmd/delete.go
@@ -68,6 +68,11 @@ var DeleteCommand = &cli.Command{
 			return err
 		}
 
+		// check for updates after running the command
+		if err := CheckForUpdate(cCtx, afero.NewOsFs()); err != nil {
+			return err
+		}
+
 		return nil
 	},
 }

diff --git a/cmd/import.go b/cmd/import.go
@@ -93,6 +93,8 @@ var ImportCommand = &cli.Command{
 		if err != nil {
 			return err
 		}
+
+		// set the number of workers based on the number of CPUs
 		numParsers = int(math.Floor(math.Max(4, float64(runtime.NumCPU())/2)))
 		numDigesters = int(math.Floor(math.Max(4, float64(runtime.NumCPU())/2)))
 		numWriters = int(math.Floor(math.Max(4, float64(runtime.NumCPU())/2)))
@@ -102,7 +104,16 @@ var ImportCommand = &cli.Command{
 
 		// run import command
 		_, err = RunImportCmd(startTime, cfg, afs, cCtx.String("logs"), cCtx.String("database"), cCtx.Bool("rolling"), cCtx.Bool("rebuild"))
-		return err
+		if err != nil {
+			return err
+		}
+
+		// check for updates after running the command
+		if err := CheckForUpdate(cCtx, afero.NewOsFs()); err != nil {
+			return err
+		}
+
+		return nil
 	},
 }
 

diff --git a/cmd/list.go b/cmd/list.go
@@ -38,6 +38,11 @@ var ListCommand = &cli.Command{
 			return err
 		}
 
+		// check for updates after running the command
+		if err := CheckForUpdate(cCtx, afero.NewOsFs()); err != nil {
+			return err
+		}
+
 		return nil
 	},
 }

diff --git a/cmd/validate.go b/cmd/validate.go
@@ -35,6 +35,12 @@ var ValidateConfigCommand = &cli.Command{
 
 		// validate config file
 		if err := RunValidateConfigCommand(afs, cCtx.String("config")); err != nil {
+			fmt.Printf("\n\t[!] Configuration file is not valid...")
+			return err
+		}
+
+		// check for updates after running the command
+		if err := CheckForUpdate(cCtx, afero.NewOsFs()); err != nil {
 			return err
 		}
 
@@ -54,7 +60,7 @@ func RunValidateConfigCommand(afs afero.Fs, configPath string) error {
 		return err
 	}
 
-	fmt.Println("configuration file is valid")
+	fmt.Printf("\n\t[✨] Configuration file is valid \n\n")
 
 	return nil
 }

diff --git a/cmd/view.go b/cmd/view.go
@@ -77,8 +77,16 @@ var ViewCommand = &cli.Command{
 			}
 		}
 
-		err := runViewCmd(afs, cCtx.String("config"), cCtx.Args().First(), cCtx.Bool("stdout"), cCtx.String("search"), cCtx.Int("limit"))
-		return err
+		if err := runViewCmd(afs, cCtx.String("config"), cCtx.Args().First(), cCtx.Bool("stdout"), cCtx.String("search"), cCtx.Int("limit")); err != nil {
+			return err
+		}
+
+		// check for updates after running the command
+		if err := CheckForUpdate(cCtx, afero.NewOsFs()); err != nil {
+			return err
+		}
+
+		return nil
 	},
 }
 

diff --git a/rita.go b/rita.go
@@ -4,13 +4,11 @@ import (
 	"activecm/rita/cmd"
 	"activecm/rita/config"
 	"activecm/rita/logger"
-	"activecm/rita/util"
 	"activecm/rita/viewer"
 	"fmt"
 	"log"
 	"os"
 
-	"github.com/google/go-github/github"
 	"github.com/joho/godotenv"
 	"github.com/urfave/cli/v2"
 )
@@ -61,17 +59,6 @@ func main() {
 				log.Fatal("Error loading .env file", err)
 			}
 
-			// check for update if version is set
-			if Version != "" {
-				newer, latestVersion, err := util.CheckForNewerVersion(github.NewClient(nil), "v0.0.0")
-				if err != nil {
-					log.Fatalf("Error checking for newer version: %v", err)
-				}
-				if newer {
-					fmt.Printf("\n\t✨ A newer version (%s) of RITA is available! https://github.com/activecm/rita/releases ✨\n", latestVersion)
-				}
-			}
-
 			return nil
 		},
 	}
@@ -88,7 +75,7 @@ func exitErrHandler(c *cli.Context, err error) {
 	if err == nil {
 		return
 	}
-	fmt.Fprintf(c.App.ErrWriter, "\n[!] %+v\n", err.Error())
+	fmt.Fprintf(c.App.ErrWriter, "\n\n\t[!] %+v\n\n", err.Error())
 	cli.OsExiter(1)
 
 }