diff --git a/cmd/cmd.go b/cmd/cmd.go index 59f2f95..a7123af 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -15,6 +15,9 @@ import ( var ErrMissingDatabaseName = errors.New("database name is required") var ErrMissingConfigPath = errors.New("config path parameter is required") var ErrTooManyArguments = errors.New("too many arguments provided") +var ErrInvalidConfigObject = errors.New("config was nil or invalid") +var ErrCurrentVersionEmpty = errors.New("current version unset") +var ErrCheckingForUpdate = errors.New("error checking for newer version of RITA") func Commands() []*cli.Command { return []*cli.Command{ @@ -39,21 +42,20 @@ func ConfigFlag(required bool) *cli.StringFlag { } } -func CheckForUpdate(cCtx *cli.Context, afs afero.Fs) error { +func CheckForUpdate(cfg *config.Config) error { + // make sure config is not nil + if cfg == nil { + return ErrInvalidConfigObject + } + // get the current version currentVersion := config.Version - // load config file - cfg, err := config.ReadFileConfig(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 != "" && currentVersion != "dev" { newer, latestVersion, err := util.CheckForNewerVersion(github.NewClient(nil), currentVersion) if err != nil { - return fmt.Errorf("error checking for newer version of RITA: %w", err) + return fmt.Errorf("%w: %w", ErrCheckingForUpdate, err) } if newer { fmt.Printf("\n\t✨ A newer version (%s) of RITA is available! https://github.com/activecm/rita/releases ✨\n\n", latestVersion) diff --git a/cmd/cmd_test.go b/cmd/cmd_test.go index 781e394..9e9d704 100644 --- a/cmd/cmd_test.go +++ b/cmd/cmd_test.go @@ -1,6 +1,7 @@ package cmd_test import ( + "bytes" "context" "fmt" "log" @@ -8,11 +9,15 @@ import ( "path/filepath" "testing" + "github.com/activecm/rita/v5/cmd" "github.com/activecm/rita/v5/config" "github.com/activecm/rita/v5/database" + "github.com/activecm/rita/v5/util" + "github.com/google/go-github/github" "github.com/joho/godotenv" "github.com/spf13/afero" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" "github.com/testcontainers/testcontainers-go" @@ -91,20 +96,7 @@ func (c *CmdTestSuite) TearDownSuite() { // func (d *DatabaseTestSuite) TearDownTest() {} // SetupSubTest is run before each subtest -func (c *CmdTestSuite) SetupSubTest() { - t := c.T() - fmt.Println("Running setup subtest...") - - // drop all databases that may have been created during subtest - if c.server != nil && c.server.Conn != nil { - dbs, err := c.server.ListImportDatabases() - require.NoError(t, err, "listing databases should not produce an error") - for _, db := range dbs { - err := c.server.DeleteSensorDB(db.Name) - require.NoError(t, err, "dropping database should not produce an error") - } - } -} +// func (c *CmdTestSuite) SetupSubTest() {} // TearDownSubTest is run after each subtest // func (c *CmdTestSuite) TearDownSubTest() {} @@ -148,8 +140,117 @@ func setupTestApp(commands []*cli.Command, flags []cli.Flag) (*cli.App, context. // this prevents the test from exiting when testing for errors app.ExitErrHandler = func(_ *cli.Context, _ error) { // add any custom test logic, or assertions or leave it blank - } return app, ctx } + +func TestCheckForUpdate(t *testing.T) { + // set up file system interface + afs := afero.NewOsFs() + + // load the config file + cfg, err := config.ReadFileConfig(afs, ConfigPath) + require.NoError(t, err, "config should load without error") + + // get latest release version + latestVersion, err := util.GetLatestReleaseVersion(github.NewClient(nil), "activecm", "rita") + require.NoError(t, err, "latest release version should be retrieved without error") + + tests := []struct { + name string + cfg *config.Config + updateCheckEnabled bool + currentVersion string + expectedErr error + expectedOutput string + }{ + { + name: "New version available", + updateCheckEnabled: true, + cfg: cfg, + currentVersion: "v0.0.0", + expectedOutput: fmt.Sprintf("\n\t✨ A newer version (%s) of RITA is available! https://github.com/activecm/rita/releases ✨\n\n", latestVersion), + }, + { + name: "Error checking for newer version", + updateCheckEnabled: true, + cfg: cfg, + currentVersion: "notaversion", + expectedErr: cmd.ErrCheckingForUpdate, + }, + { + name: "Update check disabled", + updateCheckEnabled: false, + cfg: cfg, + currentVersion: "1.0.0", + }, + { + name: "Current version is dev", + updateCheckEnabled: true, + cfg: cfg, + currentVersion: "dev", + }, + { + name: "Current version is empty", + updateCheckEnabled: true, + cfg: cfg, + currentVersion: "", + }, + { + name: "Nil config", + cfg: nil, + expectedErr: cmd.ErrInvalidConfigObject, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + // set update check enabled in config + if test.cfg != nil { + test.cfg.UpdateCheckEnabled = test.updateCheckEnabled + } + + // override global variables and functions + config.Version = test.currentVersion + + // capture stdout + output := captureOutput(t, func() { + err := cmd.CheckForUpdate(test.cfg) + // check error + if test.expectedErr != nil { + require.Contains(t, err.Error(), test.expectedErr.Error(), "error should contain expected value") + } else { + assert.NoError(t, err) + } + }) + + // Assert output + if test.expectedOutput != "" { + assert.Equal(t, test.expectedOutput, output) + } + }) + } +} + +// captureOutput captures stdout from a function +func captureOutput(t *testing.T, f func()) string { + t.Helper() + + // capture stdout + old := os.Stdout + r, w, err := os.Pipe() + require.NoError(t, err) + os.Stdout = w + + // run the function + f() + + // close and restore stdout + w.Close() + os.Stdout = old + var buf bytes.Buffer + _, err = buf.ReadFrom(r) + require.NoError(t, err) + return buf.String() +} diff --git a/cmd/delete.go b/cmd/delete.go index a55d19b..8e4ccaa 100644 --- a/cmd/delete.go +++ b/cmd/delete.go @@ -77,7 +77,7 @@ var DeleteCommand = &cli.Command{ } // check for updates after running the command - if err := CheckForUpdate(cCtx, afero.NewOsFs()); err != nil { + if err := CheckForUpdate(cfg); err != nil { return err } diff --git a/cmd/delete_test.go b/cmd/delete_test.go index f7f685b..008060d 100644 --- a/cmd/delete_test.go +++ b/cmd/delete_test.go @@ -202,8 +202,12 @@ func (c *CmdTestSuite) TestRunDeleteCmd() { } // validate that the expected databases remain - for _, db := range test.expectedRemainingDbs { - require.Contains(t, dbString, db, "database %s should not have been deleted", db) + require.ElementsMatch(t, test.expectedRemainingDbs, dbString, "remaining databases should match expected value") + + // cleanup + for _, db := range test.dbs { + err := c.server.DeleteSensorDB(db.name) + require.NoError(t, err, "dropping database should not produce an error") } }) diff --git a/cmd/import.go b/cmd/import.go index a0b3516..044d1ef 100644 --- a/cmd/import.go +++ b/cmd/import.go @@ -40,6 +40,7 @@ var ErrInvalidLogHourRange = errors.New("could not parse hour from log file name var ErrInvalidLogType = errors.New("incompatible log type") var ErrIncompatibleFileExtension = errors.New("incompatible file extension") var ErrSkippedDuplicateLog = errors.New("encountered file with same name but different extension, skipping file due to older last modified time") +var ErrMissingLogDirectory = errors.New("log directory flag is required") type WalkError struct { Path string @@ -110,7 +111,7 @@ var ImportCommand = &cli.Command{ } // check for updates after running the command - if err := CheckForUpdate(cCtx, afero.NewOsFs()); err != nil { + if err := CheckForUpdate(cfg); err != nil { return err } @@ -303,7 +304,7 @@ func RunImportCmd(startTime time.Time, cfg *config.Config, afs afero.Fs, logDir func ValidateLogDirectory(afs afero.Fs, logDir string) error { if logDir == "" { - return fmt.Errorf("log directory flag is required") + return ErrMissingLogDirectory } dir, err := util.ParseRelativePath(logDir) @@ -343,7 +344,12 @@ func ValidateDatabaseName(name string) error { return nil } -func parseFolderDate(folder string) (time.Time, error) { +// ParseFolderDate extracts the date from a given folder name +func ParseFolderDate(folder string) (time.Time, error) { + if folder == "" { + return time.Unix(0, 0), errors.New("folder name cannot be empty") + } + // check if the path is a directory folderDate, err := time.Parse(time.DateOnly, folder) if err != nil { @@ -489,7 +495,7 @@ func WalkFiles(afs afero.Fs, root string) ([]HourlyZeekLogs, []WalkError, error) } parentDir := filepath.Base(filepath.Dir(file.path)) - folderDate, err := parseFolderDate(parentDir) + folderDate, err := ParseFolderDate(parentDir) if err != nil { walkErrors = append(walkErrors, WalkError{Path: path, Error: err}) } diff --git a/cmd/import_test.go b/cmd/import_test.go index f7c7545..0114b84 100644 --- a/cmd/import_test.go +++ b/cmd/import_test.go @@ -52,56 +52,193 @@ NON-ROLLING LOGS */ func (c *CmdTestSuite) TestRunImportCmd() { - - type TestCase struct { + type importDB struct { name string - afs afero.Fs - dbName string - rolling bool - rebuild bool logDir string hours [][]string + rolling bool + rebuild bool expectedImport int + expectedError error + } + + type TestCase struct { + name string + afs afero.Fs + importDBs []importDB + // dbName string + // rolling bool + // rebuild bool + // logDir string + // hours [][]string + // expectedImport int + // expectedError error } testCases := []TestCase{ { - name: "No Subdirectories, No Hours", - afs: afero.NewOsFs(), - dbName: "ahhhhhhhhhh", - rolling: false, - rebuild: false, - logDir: "../test_data/valid_tsv", - hours: [][]string{{"conn.log.gz", "dns.log.gz", "http.log.gz", "ssl.log.gz", "open_conn.log.gz", "open_http.log.gz", "open_ssl.log.gz"}}, - expectedImport: 1, + name: "No Subdirectories, No Hours", + afs: afero.NewOsFs(), + importDBs: []importDB{ + { + name: "ahhhhhhhhhh", + logDir: "../test_data/valid_tsv", + hours: [][]string{{"conn.log.gz", "dns.log.gz", "http.log.gz", "ssl.log.gz", "open_conn.log.gz", "open_http.log.gz", "open_ssl.log.gz"}}, + rolling: false, + rebuild: false, + expectedImport: 1, + expectedError: nil, + }, + }, }, { - name: "Simple, SubDirectories - Multi-Day Logs", - afs: afero.NewMemMapFs(), - dbName: "bingbong", - rolling: false, - rebuild: false, - logDir: "/logs", - hours: [][]string{ - {"2024-04-29/conn.log", "2024-04-29/dns.log", "2024-04-29/http.log", "2024-04-29/ssl.log", "2024-04-29/open_conn.log", "2024-04-29/open_http.log", "2024-04-29/open_ssl.log"}, - {"2024-05-01/conn.log", "2024-05-01/dns.log", "2024-05-01/http.log", "2024-05-01/ssl.log", "2024-05-01/open_conn.log", "2024-05-01/open_http.log", "2024-05-01/open_ssl.log", "2024-05-01/ssl_blue.log"}, + name: "Simple, SubDirectories - Multi-Day Logs", + afs: afero.NewMemMapFs(), + importDBs: []importDB{ + { + name: "bingbong", + logDir: "/logs", + hours: [][]string{ + {"2024-04-29/conn.log", "2024-04-29/dns.log", "2024-04-29/http.log", "2024-04-29/ssl.log", "2024-04-29/open_conn.log", "2024-04-29/open_http.log", "2024-04-29/open_ssl.log"}, + {"2024-05-01/conn.log", "2024-05-01/dns.log", "2024-05-01/http.log", "2024-05-01/ssl.log", "2024-05-01/open_conn.log", "2024-05-01/open_http.log", "2024-05-01/open_ssl.log", "2024-05-01/ssl_blue.log"}, + }, + rolling: false, + rebuild: false, + expectedImport: 2, + expectedError: nil, + }, }, - expectedImport: 2, }, { - name: "SubDirectories, Multi-Day, Multi-Hour Logs", - afs: afero.NewMemMapFs(), - dbName: "bingbong", - rolling: false, - rebuild: false, - logDir: "/logs", - hours: [][]string{ - {"2024-04-29/conn.00:00:00-01:00:00.log", "2024-04-29/open_conn.00:00:00-01:00:00.log", "2024-04-29/dns.00:00:00-01:00:00.log", "2024-04-29/http.00:00:00-01:00:00.log", "2024-04-29/open_http.00:00:00-01:00:00.log", "2024-04-29/ssl.00:00:00-01:00:00.log", "2024-04-29/open_ssl.00:00:00-01:00:00.log"}, - {"2024-04-29/conn.23:00:00-00:00:00.log", "2024-04-29/open_conn.23:00:00-00:00:00.log", "2024-04-29/dns.23:00:00-00:00:00.log", "2024-04-29/http.23:00:00-00:00:00.log", "2024-04-29/open_http.23:00:00-00:00:00.log", "2024-04-29/ssl.23:00:00-00:00:00.log", "2024-04-29/open_ssl.23:00:00-00:00:00.log"}, - {"2024-05-01/conn.00:00:00-01:00:00.log", "2024-05-01/open_conn.00:00:00-01:00:00.log", "2024-05-01/dns.00:00:00-01:00:00.log", "2024-05-01/http.00:00:00-01:00:00.log", "2024-05-01/open_http.00:00:00-01:00:00.log", "2024-05-01/ssl.00:00:00-01:00:00.log", "2024-05-01/open_ssl.00:00:00-01:00:00.log"}, - {"2024-05-01/conn.23:00:00-00:00:00.log", "2024-05-01/open_conn.23:00:00-00:00:00.log", "2024-05-01/dns.23:00:00-00:00:00.log", "2024-05-01/http.23:00:00-00:00:00.log", "2024-05-01/open_http.23:00:00-00:00:00.log", "2024-05-01/ssl.23:00:00-00:00:00.log", "2024-05-01/open_ssl.23:00:00-00:00:00.log", "2024-05-01/ssl_blue.23:00:00-00:00:00.log"}, + name: "SubDirectories, Multi-Day, Multi-Hour Logs", + afs: afero.NewMemMapFs(), + importDBs: []importDB{ + { + name: "bingbong", + logDir: "/logs", + hours: [][]string{ + {"2024-04-29/conn.00:00:00-01:00:00.log", "2024-04-29/open_conn.00:00:00-01:00:00.log", "2024-04-29/dns.00:00:00-01:00:00.log", "2024-04-29/http.00:00:00-01:00:00.log", "2024-04-29/open_http.00:00:00-01:00:00.log", "2024-04-29/ssl.00:00:00-01:00:00.log", "2024-04-29/open_ssl.00:00:00-01:00:00.log"}, + {"2024-04-29/conn.23:00:00-00:00:00.log", "2024-04-29/open_conn.23:00:00-00:00:00.log", "2024-04-29/dns.23:00:00-00:00:00.log", "2024-04-29/http.23:00:00-00:00:00.log", "2024-04-29/open_http.23:00:00-00:00:00.log", "2024-04-29/ssl.23:00:00-00:00:00.log", "2024-04-29/open_ssl.23:00:00-00:00:00.log"}, + {"2024-05-01/conn.00:00:00-01:00:00.log", "2024-05-01/open_conn.00:00:00-01:00:00.log", "2024-05-01/dns.00:00:00-01:00:00.log", "2024-05-01/http.00:00:00-01:00:00.log", "2024-05-01/open_http.00:00:00-01:00:00.log", "2024-05-01/ssl.00:00:00-01:00:00.log", "2024-05-01/open_ssl.00:00:00-01:00:00.log"}, + {"2024-05-01/conn.23:00:00-00:00:00.log", "2024-05-01/open_conn.23:00:00-00:00:00.log", "2024-05-01/dns.23:00:00-00:00:00.log", "2024-05-01/http.23:00:00-00:00:00.log", "2024-05-01/open_http.23:00:00-00:00:00.log", "2024-05-01/ssl.23:00:00-00:00:00.log", "2024-05-01/open_ssl.23:00:00-00:00:00.log", "2024-05-01/ssl_blue.23:00:00-00:00:00.log"}, + }, + rolling: false, + rebuild: false, + expectedImport: 4, + expectedError: nil, + }, + }, + }, + + { + name: "SubDirectories, Multi-Day, Multi-Hour Logs", + afs: afero.NewMemMapFs(), + importDBs: []importDB{ + { + name: "bingbong", + logDir: "/logs", + hours: [][]string{ + {"2024-04-29/conn.00:00:00-01:00:00.log", "2024-04-29/open_conn.00:00:00-01:00:00.log", "2024-04-29/dns.00:00:00-01:00:00.log", "2024-04-29/http.00:00:00-01:00:00.log", "2024-04-29/open_http.00:00:00-01:00:00.log", "2024-04-29/ssl.00:00:00-01:00:00.log", "2024-04-29/open_ssl.00:00:00-01:00:00.log"}, + {"2024-04-29/conn.23:00:00-00:00:00.log", "2024-04-29/open_conn.23:00:00-00:00:00.log", "2024-04-29/dns.23:00:00-00:00:00.log", "2024-04-29/http.23:00:00-00:00:00.log", "2024-04-29/open_http.23:00:00-00:00:00.log", "2024-04-29/ssl.23:00:00-00:00:00.log", "2024-04-29/open_ssl.23:00:00-00:00:00.log"}, + {"2024-05-01/conn.00:00:00-01:00:00.log", "2024-05-01/open_conn.00:00:00-01:00:00.log", "2024-05-01/dns.00:00:00-01:00:00.log", "2024-05-01/http.00:00:00-01:00:00.log", "2024-05-01/open_http.00:00:00-01:00:00.log", "2024-05-01/ssl.00:00:00-01:00:00.log", "2024-05-01/open_ssl.00:00:00-01:00:00.log"}, + {"2024-05-01/conn.23:00:00-00:00:00.log", "2024-05-01/open_conn.23:00:00-00:00:00.log", "2024-05-01/dns.23:00:00-00:00:00.log", "2024-05-01/http.23:00:00-00:00:00.log", "2024-05-01/open_http.23:00:00-00:00:00.log", "2024-05-01/ssl.23:00:00-00:00:00.log", "2024-05-01/open_ssl.23:00:00-00:00:00.log", "2024-05-01/ssl_blue.23:00:00-00:00:00.log"}, + }, + rolling: false, + rebuild: false, + expectedImport: 4, + expectedError: nil, + }, + }, + }, + { + name: "Rolling", + afs: afero.NewMemMapFs(), + importDBs: []importDB{ + { + name: "bingbong", + logDir: "/logs/1", + hours: [][]string{ + {"conn.log", "dns.log", "http.log", "ssl.log", "open_conn.log", "open_http.log", "open_ssl.log"}, + }, + rolling: true, + rebuild: false, + expectedImport: 1, + expectedError: nil, + }, + { + name: "bingbong", + logDir: "/logs/2", + hours: [][]string{ + {"conn.log", "dns.log", "http.log", "ssl.log", "open_conn.log", "open_http.log", "open_ssl.log"}, + }, + rolling: true, + rebuild: false, + expectedImport: 1, + expectedError: nil, + }, + }, + }, + { + name: "Rolling - Multi-Day, Multi-Hour Logs", + afs: afero.NewMemMapFs(), + importDBs: []importDB{ + { + name: "bingbong", + logDir: "/logs/1", + hours: [][]string{ + {"2024-04-29/conn.00:00:00-01:00:00.log", "2024-04-29/open_conn.00:00:00-01:00:00.log", "2024-04-29/dns.00:00:00-01:00:00.log", "2024-04-29/http.00:00:00-01:00:00.log", "2024-04-29/open_http.00:00:00-01:00:00.log", "2024-04-29/ssl.00:00:00-01:00:00.log", "2024-04-29/open_ssl.00:00:00-01:00:00.log"}, + {"2024-04-29/conn.23:00:00-00:00:00.log", "2024-04-29/open_conn.23:00:00-00:00:00.log", "2024-04-29/dns.23:00:00-00:00:00.log", "2024-04-29/http.23:00:00-00:00:00.log", "2024-04-29/open_http.23:00:00-00:00:00.log", "2024-04-29/ssl.23:00:00-00:00:00.log", "2024-04-29/open_ssl.23:00:00-00:00:00.log"}, + {"2024-05-01/conn.00:00:00-01:00:00.log", "2024-05-01/open_conn.00:00:00-01:00:00.log", "2024-05-01/dns.00:00:00-01:00:00.log", "2024-05-01/http.00:00:00-01:00:00.log", "2024-05-01/open_http.00:00:00-01:00:00.log", "2024-05-01/ssl.00:00:00-01:00:00.log", "2024-05-01/open_ssl.00:00:00-01:00:00.log"}, + {"2024-05-01/conn.23:00:00-00:00:00.log", "2024-05-01/open_conn.23:00:00-00:00:00.log", "2024-05-01/dns.23:00:00-00:00:00.log", "2024-05-01/http.23:00:00-00:00:00.log", "2024-05-01/open_http.23:00:00-00:00:00.log", "2024-05-01/ssl.23:00:00-00:00:00.log", "2024-05-01/open_ssl.23:00:00-00:00:00.log", "2024-05-01/ssl_blue.23:00:00-00:00:00.log"}, + }, + rolling: true, + rebuild: false, + expectedImport: 4, + expectedError: nil, + }, + { + name: "bingbong", + logDir: "/logs/2", + hours: [][]string{ + {"2024-05-02/conn.00:00:00-01:00:00.log", "2024-05-02/open_conn.00:00:00-01:00:00.log", "2024-05-02/dns.00:00:00-01:00:00.log", "2024-05-02/http.00:00:00-01:00:00.log", "2024-05-02/open_http.00:00:00-01:00:00.log", "2024-05-02/ssl.00:00:00-01:00:00.log", "2024-05-02/open_ssl.00:00:00-01:00:00.log"}, + {"2024-05-02/conn.23:00:00-00:00:00.log", "2024-05-02/open_conn.23:00:00-00:00:00.log", "2024-05-02/dns.23:00:00-00:00:00.log", "2024-05-02/http.23:00:00-00:00:00.log", "2024-05-02/open_http.23:00:00-00:00:00.log", "2024-05-02/ssl.23:00:00-00:00:00.log", "2024-05-02/open_ssl.23:00:00-00:00:00.log"}, + {"2024-05-03/conn.00:00:00-01:00:00.log", "2024-05-03/open_conn.00:00:00-01:00:00.log", "2024-05-03/dns.00:00:00-01:00:00.log", "2024-05-03/http.00:00:00-01:00:00.log", "2024-05-03/open_http.00:00:00-01:00:00.log", "2024-05-03/ssl.00:00:00-01:00:00.log", "2024-05-03/open_ssl.00:00:00-01:00:00.log"}, + {"2024-05-03/conn.23:00:00-00:00:00.log", "2024-05-03/open_conn.23:00:00-00:00:00.log", "2024-05-03/dns.23:00:00-00:00:00.log", "2024-05-03/http.23:00:00-00:00:00.log", "2024-05-03/open_http.23:00:00-00:00:00.log", "2024-05-03/ssl.23:00:00-00:00:00.log", "2024-05-03/open_ssl.23:00:00-00:00:00.log", "2024-05-03/ssl_blue.23:00:00-00:00:00.log"}, + }, + rolling: true, + rebuild: false, + expectedImport: 4, + expectedError: nil, + }, + }, + }, + { + name: "Files Previously Imported", + afs: afero.NewMemMapFs(), + importDBs: []importDB{ + { + name: "bingbong", + logDir: "/logs/1", + hours: [][]string{ + {"conn.log", "dns.log", "http.log", "ssl.log", "open_conn.log", "open_http.log", "open_ssl.log"}, + }, + rolling: true, + rebuild: false, + expectedImport: 1, + expectedError: nil, + }, + { + name: "bingbong", + logDir: "/logs/1", + hours: [][]string{ + {"conn.log", "dns.log", "http.log", "ssl.log", "open_conn.log", "open_http.log", "open_ssl.log"}, + }, + rolling: true, + rebuild: false, + expectedImport: 1, + expectedError: importer.ErrAllFilesPreviouslyImported, + }, }, - expectedImport: 4, }, } @@ -109,88 +246,111 @@ func (c *CmdTestSuite) TestRunImportCmd() { c.Run(tc.name, func() { t := c.T() - importStartedAt := time.Now() + // loop over each importDB + for _, db := range tc.importDBs { + // get start time + importStartedAt := time.Now() - var files []string - var fullPathHours [][]string + var files []string + var fullPathHours [][]string - // get the root directory path - fullRootDir := tc.logDir + // get the root directory path + fullRootDir := db.logDir - // if we are using the real logs directory, we need to get the real full path - if tc.logDir != "/logs" { - // get the current working directory - cwd, err := os.Getwd() - if err != nil { - fmt.Println("Error getting current working directory:", err) - return + // if we are using the real logs directory, we need to get the real full path + if !strings.HasPrefix(db.logDir, "/logs") { + // get the current working directory + cwd, err := os.Getwd() + if err != nil { + fmt.Println("Error getting current working directory:", err) + return + } + + fullRootDir = filepath.Join(cwd, db.logDir) } - fullRootDir = filepath.Join(cwd, tc.logDir) - } + // iterate over each day of logs + for _, day := range db.hours { + // append all the files to a single list for creation + files = append(files, day...) - // iterate over each day of logs - for _, day := range tc.hours { - // append all the files to a single list for creation - files = append(files, day...) - - // convert the day list of files to full path versions - var fullHourFiles []string - for _, file := range day { - fullPath := filepath.Join(fullRootDir, file) - fullHourFiles = append(fullHourFiles, fullPath) - } - fullPathHours = append(fullPathHours, fullHourFiles) - } - - // if we are using the mock directory, we need to create it along with the files - if tc.logDir == "/logs" { - // create mock directory with files - createMockZeekLogs(t, tc.afs, tc.logDir, files, true) - } - - // run the import command - importResults, err := cmd.RunImportCmd(importStartedAt, c.cfg, tc.afs, tc.logDir, tc.dbName, tc.rolling, tc.rebuild) - require.NoError(t, err, "running import command should not produce an error") - require.NotNil(t, importResults, "import results should not be nil") - - // verify the number of import IDs - require.Len(t, importResults.ImportID, tc.expectedImport, "import results should have expected number of import IDs") - - // check if the database exists - exists, err := database.SensorDatabaseExists(context.Background(), c.server.Conn, tc.dbName) - require.NoError(t, err, "checking if sensor database exists should not produce an error") - require.True(t, exists, "sensor database should exist") - - // check rolling status - isRolling, err := database.GetRollingStatus(context.Background(), c.server.Conn, tc.dbName) - require.NoError(t, err, "checking if sensor database is rolling should not produce an error") - require.Equal(t, tc.rolling, isRolling, "rolling status should match expected value") - - // verify imported paths for each hour - for i := range fullPathHours { - var result struct { - Paths []string `ch:"paths"` + // convert the day list of files to full path versions + var fullHourFiles []string + for _, file := range day { + fullPath := filepath.Join(fullRootDir, file) + fullHourFiles = append(fullHourFiles, fullPath) + } + fullPathHours = append(fullPathHours, fullHourFiles) } - ctx := clickhouse.Context(context.Background(), clickhouse.WithParameters(clickhouse.Parameters{ - "import_id": importResults.ImportID[i].Hex(), - "database": tc.dbName, - })) + // if we are using the mock directory, we need to create it along with the files + if strings.HasPrefix(db.logDir, "/logs") { + // create mock directory with files + createMockZeekLogs(t, tc.afs, db.logDir, files, true) + } - err = c.server.Conn.QueryRow(ctx, ` + // run the import command + importResults, err := cmd.RunImportCmd(importStartedAt, c.cfg, tc.afs, db.logDir, db.name, db.rolling, db.rebuild) + + // check if we expect an error + if db.expectedError != nil { + require.Error(t, err, "running import command should produce an error") + require.Contains(t, err.Error(), db.expectedError.Error(), "error should contain expected value") + continue + } + + // if no error was expected, continue with the rest of the checks + require.NoError(t, err, "running import command should not produce an error") + require.NotNil(t, importResults, "import results should not be nil") + + // verify the number of import IDs + require.Len(t, importResults.ImportID, db.expectedImport, "import results should have expected number of import IDs") + + // check if the database exists + exists, err := database.SensorDatabaseExists(context.Background(), c.server.Conn, db.name) + require.NoError(t, err, "checking if sensor database exists should not produce an error") + require.True(t, exists, "sensor database should exist") + + // check rolling status + isRolling, err := database.GetRollingStatus(context.Background(), c.server.Conn, db.name) + require.NoError(t, err, "checking if sensor database is rolling should not produce an error") + require.Equal(t, db.rolling, isRolling, "rolling status should match expected value") + + // verify imported paths for each hour + for i := range fullPathHours { + var result struct { + Paths []string `ch:"paths"` + } + + ctx := clickhouse.Context(context.Background(), clickhouse.WithParameters(clickhouse.Parameters{ + "import_id": importResults.ImportID[i].Hex(), + "database": db.name, + })) + + err = c.server.Conn.QueryRow(ctx, ` SELECT groupArray(path) AS paths FROM metadatabase.files WHERE import_id = unhex({import_id:String}) AND database = {database:String} `).ScanStruct(&result) - require.NoError(t, err, "querying for total file count should not produce an error") + require.NoError(t, err, "querying for total file count should not produce an error") + + require.ElementsMatch(t, fullPathHours[i], result.Paths, "paths should match expected value") + } - require.ElementsMatch(t, fullPathHours[i], result.Paths, "paths should match expected value") } - // clean up the directory - if tc.logDir == "/logs" { - require.NoError(t, tc.afs.RemoveAll(tc.logDir), "removing directory should not produce an error") + // cleanup each importDB + for _, db := range tc.importDBs { + // clean up the directory if we are using a mock directory + // if tc.logDir == "/logs" { + if strings.HasPrefix(db.logDir, "/logs") { + require.NoError(t, tc.afs.RemoveAll(db.logDir), "removing directory should not produce an error") + } + + // clean up the database + err := c.server.DeleteSensorDB(db.name) + require.NoError(t, err, "dropping database should not produce an error") + } }) } @@ -230,7 +390,7 @@ func createMockZeekLogs(t *testing.T, afs afero.Fs, directory string, files []st "1715641234.367201\tCxT125\t10.0.0.5\t52.12.0.5\n", ) } - err = afero.WriteFile(afs, filepath.Join(directory, file), data, os.FileMode(0o775)) + err := afero.WriteFile(afs, filepath.Join(directory, file), data, os.FileMode(0o775)) require.NoError(t, err, "creating files should not produce an error") } } @@ -949,21 +1109,25 @@ func TestValidateDatabaseName(t *testing.T) { } tests := []testCase{ + {name: "Common name, dnscat2_ja3_strobe", db: "dnscat2_ja3_strobe"}, + {name: "Common name, combined__0000_rolling", db: "combined__0000_rolling"}, + {name: "Common name, seconion_2024_05_15", db: "combined__0000_rolling"}, {name: "All alpha characters", db: "vsagent"}, {name: "All alphanumeric characters", db: "dnscat20"}, + {name: "All numeric characters", db: "2024", shouldErr: true}, + {name: "Starting with a number", db: "2vsagent", shouldErr: true}, {name: "Starting with a capital letter", db: "Vsagent", shouldErr: true}, {name: "All caps", db: "INFORMATION_SCHEMA", shouldErr: true}, + {name: "Contains special characters", db: "ch!ck3n$tr!p", shouldErr: true}, + {name: "Contains a hyphen", db: "combined__0000-rolling", shouldErr: true}, {name: "Starting with an underscore", db: "_vsagent", shouldErr: true}, + {name: "Ending with underscore", db: "dnscat2_", shouldErr: true}, + {name: "Length >63 characters", db: "i_am_a_very_long_database_name_that_is_over_63_characters_long_and_should_fail", shouldErr: true}, {name: "Name is reserved: default", db: "default", shouldErr: true}, {name: "Name is reserved: system", db: "system", shouldErr: true}, {name: "Name is reserved: information_schema", db: "information_schema", shouldErr: true}, {name: "Name is reserved: metadatabase", db: "metadatabase", shouldErr: true}, - {name: "Contains special characters", db: "ch!ck3n$tr!p", shouldErr: true}, - {name: "Contains a hyphen", db: "combined__0000-rolling", shouldErr: true}, - {name: "Ends with underscore", db: "dnscat2_", shouldErr: true}, - {name: "Common name, dnscat2_ja3_strobe", db: "dnscat2_ja3_strobe"}, - {name: "Common name, combined__0000_rolling", db: "combined__0000_rolling"}, - {name: "Common name, seconion_2024_05_15", db: "combined__0000_rolling"}, + {name: "Empty string", db: "", shouldErr: true}, } for _, test := range tests { @@ -973,3 +1137,114 @@ func TestValidateDatabaseName(t *testing.T) { }) } } + +func TestValidateLogDirectory(t *testing.T) { + tests := []struct { + name string + logDir string + setup func(afs afero.Fs) + expectedError error + }{ + { + name: "Valid Directory", + logDir: "/validlogdir", + setup: func(afs afero.Fs) { + require.NoError(t, afs.Mkdir("/validlogdir", 0755)) + require.NoError(t, afero.WriteFile(afs, "/validlogdir/file.txt", []byte("content"), 0644)) + }, + expectedError: nil, + }, + { + name: "Empty Directory", + logDir: "/emptylogdir", + setup: func(afs afero.Fs) { + require.NoError(t, afs.Mkdir("/emptylogdir", 0755)) + }, + expectedError: util.ErrDirIsEmpty, + }, + { + name: "Path is a File", + logDir: "/logfile.txt", + setup: func(afs afero.Fs) { + require.NoError(t, afero.WriteFile(afs, "/logfile.txt", []byte("content"), 0644)) + }, + expectedError: util.ErrPathIsNotDir, + }, + { + name: "Empty Log Directory", + logDir: "", + setup: func(_ afero.Fs) {}, + expectedError: cmd.ErrMissingLogDirectory, + }, + { + name: "Invalid Relative Path", + logDir: "~/invalid/dir", + setup: func(_ afero.Fs) {}, + expectedError: util.ErrDirDoesNotExist, + }, + { + name: "Non-Existent Directory", + logDir: "/nonexistentdir", + setup: func(_ afero.Fs) {}, + expectedError: util.ErrDirDoesNotExist, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + afs := afero.NewMemMapFs() + test.setup(afs) + + err := cmd.ValidateLogDirectory(afs, test.logDir) + + if test.expectedError != nil { + require.Error(t, err, "error should not be nil") + require.ErrorContains(t, err, test.expectedError.Error(), "error message should contain expected value") + } else { + require.NoError(t, err, "validating log directory should not produce an error") + } + }) + } +} + +func TestParseFolderDate(t *testing.T) { + tests := []struct { + name string + folder string + expectedTime time.Time + expectedError error + }{ + { + name: "Valid Date Folder", + folder: "2023-06-01", + expectedTime: time.Date(2023, 6, 1, 0, 0, 0, 0, time.UTC), + }, + { + name: "Invalid Date Folder", + folder: "invalid-folder", + expectedTime: time.Date(2006, 1, 2, 0, 0, 0, 0, time.UTC), + }, + { + name: "Empty Folder Name", + folder: "", + expectedTime: time.Unix(0, 0), + expectedError: fmt.Errorf("folder name cannot be empty"), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result, err := cmd.ParseFolderDate(test.folder) + + if test.expectedError != nil { + require.Error(t, err, "error should not be nil") + require.ErrorContains(t, err, test.expectedError.Error(), "error message should contain expected value") + + } else { + require.NoError(t, err, "parsing folder date should not produce an error") + } + + require.Equal(t, test.expectedTime, result, "the result should match the expected value") + }) + } +} diff --git a/cmd/list.go b/cmd/list.go index 691e15d..ac74652 100644 --- a/cmd/list.go +++ b/cmd/list.go @@ -34,13 +34,19 @@ var ListCommand = &cli.Command{ // set up file system interface afs := afero.NewOsFs() + // load config file + cfg, err := config.ReadFileConfig(afs, cCtx.String("config")) + if err != nil { + return err + } + // run the delete command - if err := runListCmd(afs, cCtx.String("config")); err != nil { + if err := runListCmd(cfg); err != nil { return err } // check for updates after running the command - if err := CheckForUpdate(cCtx, afero.NewOsFs()); err != nil { + if err := CheckForUpdate(cfg); err != nil { return err } @@ -48,12 +54,7 @@ var ListCommand = &cli.Command{ }, } -func runListCmd(afs afero.Fs, configPath string) error { - - cfg, err := config.ReadFileConfig(afs, configPath) - if err != nil { - return err - } +func runListCmd(cfg *config.Config) error { // connect to server server, err := database.ConnectToServer(context.Background(), cfg) diff --git a/cmd/list_test.go b/cmd/list_test.go index 8f27755..6229fc8 100644 --- a/cmd/list_test.go +++ b/cmd/list_test.go @@ -66,4 +66,10 @@ func (c *CmdTestSuite) TestFormatListTable() { require.Equal(expectedDBs[i].rolling, strings.TrimSpace(cols[1])) require.Equal(expectedDBs[i].tsRange, strings.TrimSpace(cols[2])) } + + // clean up + for _, db := range expectedDBs { + err := c.server.DeleteSensorDB(db.name) + require.NoError(err) + } } diff --git a/cmd/validate.go b/cmd/validate.go index d8e1b1d..9b6caf6 100644 --- a/cmd/validate.go +++ b/cmd/validate.go @@ -35,13 +35,14 @@ var ValidateConfigCommand = &cli.Command{ afs := afero.NewOsFs() // validate config file - if err := RunValidateConfigCommand(afs, cCtx.String("config")); err != nil { + cfg, err := RunValidateConfigCommand(afs, cCtx.String("config")) + if 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 { + if err := CheckForUpdate(cfg); err != nil { return err } @@ -49,21 +50,21 @@ var ValidateConfigCommand = &cli.Command{ }, } -func RunValidateConfigCommand(afs afero.Fs, configPath string) error { +func RunValidateConfigCommand(afs afero.Fs, configPath string) (*config.Config, error) { // validate config file path if err := ValidateConfigPath(afs, configPath); err != nil { - return err + return nil, err } // load config path - _, err := config.ReadFileConfig(afs, configPath) + cfg, err := config.ReadFileConfig(afs, configPath) if err != nil { - return err + return nil, err } fmt.Printf("\n\t[✨] Configuration file is valid \n\n") - return nil + return cfg, nil } func ValidateConfigPath(afs afero.Fs, configPath string) error { diff --git a/cmd/view.go b/cmd/view.go index 18d0010..9d2f8f2 100644 --- a/cmd/view.go +++ b/cmd/view.go @@ -46,8 +46,6 @@ var ViewCommand = &cli.Command{ ConfigFlag(false), }, Action: func(cCtx *cli.Context) error { - afs := afero.NewOsFs() - // flags must go before the argument, otherwise they won't be applied ._. // we can either make the db name a flag or see if cobra is any better if !cCtx.Args().Present() { @@ -79,18 +77,22 @@ var ViewCommand = &cli.Command{ } } + // set up file system interface + afs := afero.NewOsFs() + // load config file cfg, err := config.ReadFileConfig(afs, cCtx.String("config")) if err != nil { return err } + // run the view command if err := runViewCmd(cfg, 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 { + if err := CheckForUpdate(cfg); err != nil { return err } diff --git a/database/server.go b/database/server.go index 6a928a0..f575691 100644 --- a/database/server.go +++ b/database/server.go @@ -26,7 +26,7 @@ var ErrNoMetaDBImportRecordForDatabase = errors.New("no import record found for var ErrDatabaseNotFound = errors.New("database does not exist") var ErrDatabaseNameEmpty = errors.New("database name cannot be empty") var ErrMissingConfig = errors.New("config cannot be nil") -var errImportTwiceNonRolling = errors.New("cannot import more than once to a non-rolling database") +var ErrImportTwiceNonRolling = errors.New("cannot import more than once to a non-rolling database") var errRollingStatusFailure = errors.New("failed to detect rolling status of given import database") var errRollingFlagMissing = errors.New("cannot import non-rolling data to a rolling database") @@ -337,8 +337,8 @@ func (server *ServerConn) checkRolling(dbName string, rollingFlag bool, rebuildF // command is requesting to import data as rolling, but dataset is not rolling case rollingFlag && !rolling && !rebuildFlag: logger.Warn().Str("database", dbName). - Msg(errImportTwiceNonRolling.Error()) - return rolling, errImportTwiceNonRolling + Msg(ErrImportTwiceNonRolling.Error()) + return rolling, ErrImportTwiceNonRolling // command is requesting to import data as non-rolling, but dataset is rolling case rolling && !rollingFlag && !rebuildFlag: diff --git a/util/util.go b/util/util.go index 770f5dc..15d74da 100644 --- a/util/util.go +++ b/util/util.go @@ -66,6 +66,7 @@ func init() { privateIPBlocks = privateIPs } +// NewFixedStringHash creates a FixedString from a hash of all the passed in strings func NewFixedStringHash(args ...string) (FixedString, error) { if len(args) == 0 { return FixedString{}, errors.New("no arguments provided") @@ -85,26 +86,23 @@ func NewFixedStringHash(args ...string) (FixedString, error) { return fs, nil } +// NewFixedStringFromString creates a FixedString from a passed in hex string func NewFixedStringFromHex(h string) (FixedString, error) { + if h == "" { + return FixedString{}, errors.New("hex string is empty") + } + data, err := hex.DecodeString(h) if err != nil { - return FixedString{}, err + return FixedString{}, fmt.Errorf("error decoding hex string: %w", err) } var fixed [16]byte - copy(fixed[:], data[:16]) + copy(fixed[:], data) return FixedString{ Data: fixed, }, nil } -func ValidFQDN(value string) bool { - // Regular expression for validating FQDN - // This pattern requires at least two labels (separated by dots), with each label starting and ending with an alphanumeric character. - // Labels in between can have hyphens. The last label (TLD) must be at least two characters long, with only letters. - re := regexp.MustCompile(`^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}$`) - return re.MatchString(value) -} - func (bin *FixedString) Hex() string { return strings.ToUpper(hex.EncodeToString(bin.Data[:])) } @@ -127,6 +125,14 @@ func (bin FixedString) Value() (driver.Value, error) { return &bin.val, nil } +func ValidFQDN(value string) bool { + // Regular expression for validating FQDN + // This pattern requires at least two labels (separated by dots), with each label starting and ending with an alphanumeric character. + // Labels in between can have hyphens. The last label (TLD) must be at least two characters long, with only letters. + re := regexp.MustCompile(`^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}$`) + return re.MatchString(value) +} + // ContainsIP checks if a collection of subnets contains an IP func ContainsIP(subnets []*net.IPNet, ip net.IP) bool { // cache IPv4 conversion so it not performed every in every Contains call @@ -294,23 +300,7 @@ func GetRelativeFirstSeenTimestamp(useCurrentTime bool, maxTimestamp time.Time) return time.Now() } -// func ParseRelativePath(dir string) (string, error) { -// // if path is home, parse and set home dir -// if dir[:2] == "~/" { -// home, err := os.UserHomeDir() -// if err != nil { -// return "", err -// } -// return filepath.Join(home, dir[2:]), nil -// } -// // otherwise, get the path relative to the current working directory -// currentDir, err := os.Getwd() -// if err != nil { -// return "", err -// } -// return filepath.Join(currentDir, dir), nil -// } - +// ParseRelativePath parses a given directory path and returns the absolute path func ParseRelativePath(dir string) (string, error) { // validate parameters if dir == "" { @@ -337,7 +327,6 @@ func ParseRelativePath(dir string) (string, error) { return dir, nil } - } // ValidateDirectory returns whether a directory exists and is empty @@ -392,6 +381,7 @@ func ValidateFile(afs afero.Fs, file string) error { return nil } +// validatePath validates a given path func validatePath(afs afero.Fs, path string) (bool, bool, bool, error) { var exists, isDir, isEmpty bool @@ -428,31 +418,42 @@ func validatePath(afs afero.Fs, path string) (bool, bool, bool, error) { // CheckForNewerVersion checks if a newer version of the project is available on the GitHub repository func CheckForNewerVersion(client *github.Client, currentVersion string) (bool, string, error) { - - // Get the latest release - latestRelease, _, err := client.Repositories.GetLatestRelease(context.Background(), "activecm", "rita") + // get the latest version + latestVersion, err := GetLatestReleaseVersion(client, "activecm", "rita") if err != nil { - return false, "", fmt.Errorf("error fetching latest release: %w", err) + return false, "", err } - // Get the latest version from release tag name - latestVersion := latestRelease.GetTagName() - - // Parse the current and latest versions + // parse the current version currentSemver, err := semver.ParseTolerant(currentVersion) if err != nil { return false, "", fmt.Errorf("error parsing current version: %w", err) } + // parse the latest version latestSemver, err := semver.ParseTolerant(latestVersion) if err != nil { return false, "", fmt.Errorf("error parsing latest version: %w", err) } - // Compare the versions + // compare the versions if latestSemver.GT(currentSemver) { return true, latestVersion, nil } return false, latestVersion, nil } + +// GetLatestReleaseVersion gets the latest release version from the GitHub repository +func GetLatestReleaseVersion(client *github.Client, owner, repo string) (string, error) { + // get the latest release + latestRelease, _, err := client.Repositories.GetLatestRelease(context.Background(), owner, repo) + if err != nil { + return "", fmt.Errorf("error fetching latest release: %w", err) + } + + // get the latest version from release tag name + latestVersion := latestRelease.GetTagName() + + return latestVersion, nil +} diff --git a/util/util_test.go b/util/util_test.go index b1943ff..a6062dc 100644 --- a/util/util_test.go +++ b/util/util_test.go @@ -2,6 +2,7 @@ package util import ( "crypto/md5" // #nosec G501 + "database/sql/driver" "fmt" "math" "net" @@ -14,6 +15,7 @@ import ( "github.com/google/go-github/github" "github.com/google/uuid" + "github.com/spf13/afero" "github.com/stretchr/testify/require" ) @@ -95,6 +97,240 @@ func TestNewFixedStringHash(t *testing.T) { } } +func TestNewFixedStringFromHex(t *testing.T) { + tests := []struct { + name string + input string + expected FixedString + expectedError error + }{ + { + name: "Valid Hex String", + input: "00112233445566778899aabbccddeeff", + expected: FixedString{ + Data: [16]byte{0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff}, + }, + expectedError: nil, + }, + { + name: "Valid Hex String Shorter than 16 bytes", + input: "0011223344556677", + expected: FixedString{ + Data: [16]byte{0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}, + }, + expectedError: nil, + }, + { + name: "Valid Hex String Longer than 16 bytes", + input: "00112233445566778899aabbccddeeffaabbccddeeff", + expected: FixedString{ + Data: [16]byte{0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff}, + }, + expectedError: nil, + }, + { + name: "Invalid Hex String", + input: "invalidhexstring", + expected: FixedString{}, + expectedError: fmt.Errorf("error decoding hex string: "), + }, + { + name: "Empty Hex String", + input: "", + expected: FixedString{}, + expectedError: fmt.Errorf("hex string is empty"), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result, err := NewFixedStringFromHex(test.input) + + if test.expectedError != nil { + require.Error(t, err, "error should not be nil") + require.ErrorContains(t, err, test.expectedError.Error(), "error should contain expected value") + } else { + require.NoError(t, err, "converting hex to fixed string should not produce an error") + require.Equal(t, test.expected, result, "the result should match the expected value") + } + }) + } +} + +func TestFixedString_Hex(t *testing.T) { + tests := []struct { + name string + input FixedString + expected string + }{ + { + name: "All Zeros", + input: FixedString{Data: [16]byte{}}, + expected: "00000000000000000000000000000000", + }, + { + name: "Mixed Data", + input: FixedString{Data: [16]byte{0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA, 0xB, 0xC, 0xD, 0xE, 0xF}}, + expected: "000102030405060708090A0B0C0D0E0F", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result := test.input.Hex() + require.Equal(t, test.expected, result) + }) + } +} + +func TestFixedString_MarshalBinary(t *testing.T) { + tests := []struct { + name string + input FixedString + expected []byte + }{ + { + name: "All Zeros", + input: FixedString{Data: [16]byte{}}, + expected: make([]byte, 16), + }, + { + name: "Mixed Data", + input: FixedString{Data: [16]byte{0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA, 0xB, 0xC, 0xD, 0xE, 0xF}}, + expected: []byte{0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA, 0xB, 0xC, 0xD, 0xE, 0xF}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result, err := test.input.MarshalBinary() + require.NoError(t, err) + require.Equal(t, test.expected, result) + }) + } +} + +func TestFixedString_UnmarshalBinary(t *testing.T) { + tests := []struct { + name string + input []byte + expected FixedString + }{ + { + name: "All Zeros", + input: make([]byte, 16), + expected: FixedString{Data: [16]byte{}}, + }, + { + name: "Mixed Data", + input: []byte{0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA, 0xB, 0xC, 0xD, 0xE, 0xF}, + expected: FixedString{Data: [16]byte{0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA, 0xB, 0xC, 0xD, 0xE, 0xF}}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var result FixedString + err := result.UnmarshalBinary(test.input) + require.NoError(t, err) + require.Equal(t, test.expected, result) + }) + } +} + +func TestFixedString_Value(t *testing.T) { + tests := []struct { + name string + input FixedString + expected driver.Value + }{ + { + name: "Default Value", + input: FixedString{}, + expected: "", + }, + { + name: "With Value", + input: FixedString{val: "example"}, + expected: "example", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result, err := test.input.Value() + require.NoError(t, err) + require.Equal(t, test.expected, *result.(*string)) + }) + } +} + +func TestValidFQDN(t *testing.T) { + tests := []struct { + name string + value string + expected bool + }{ + { + name: "Valid FQDN", + value: "example.com", + expected: true, + }, + { + name: "Valid FQDN with Multiple Subdomains", + value: "sub.example.com", + expected: true, + }, + { + name: "Valid FQDN with Hyphen", + value: "sub-domain.example.com", + expected: true, + }, + { + name: "Single Label", + value: "example", + expected: false, + }, + { + name: "Trailing Dot", + value: "example.com.", + expected: false, + }, + { + name: "Invalid Underscore", + value: "sub_domain.example.com", + expected: false, + }, + { + name: "Invalid Spaces", + value: "example .com", + expected: false, + }, + { + name: "Invalid Special Characters", + value: "exa$mple.com", + expected: false, + }, + { + name: "TLD Too Short", + value: "example.c", + expected: false, + }, + { + name: "Empty String", + value: "", + expected: false, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result := ValidFQDN(test.value) + require.Equal(t, test.expected, result, "the result should match the expected value") + }) + } +} + func TestContainsIP(t *testing.T) { tests := []struct { name string @@ -433,54 +669,6 @@ func TestParseNetworkID(t *testing.T) { } } -func TestValidateTimestamp(t *testing.T) { - tests := []struct { - name string - timestamp time.Time - expectedTime time.Time - replaced bool - }{ - { - name: "Valid timestamp", - timestamp: time.Date(2024, time.June, 3, 23, 24, 10, 0, time.Local), - expectedTime: time.Date(2024, time.June, 3, 23, 24, 10, 0, time.Local), - replaced: false, - }, - { - name: "Log Floating-Pont Timestamp", - timestamp: time.Unix(1517336108, int64((0.231879)*1e9)), // 1517336108.231879 - expectedTime: time.Unix(1517336108, 231879000), - replaced: false, - }, - { - name: "Unset Timestamp", - timestamp: time.Time{}, - expectedTime: time.Unix(0, 1), - replaced: true, - }, - { - name: "MaxInt64 timestamp", - timestamp: time.Unix(math.MaxInt64, 0), - expectedTime: time.Unix(0, 1), - replaced: true, - }, - { - name: "Negative timestamp", - timestamp: time.Unix(-1, 0), - expectedTime: time.Unix(0, 1), - replaced: true, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - ts, replaced := ValidateTimestamp(test.timestamp) - require.Equal(t, test.expectedTime, ts, "timestamp should match expected value") - require.Equal(t, test.replaced, replaced, "replaced should match expected value") - }) - } -} - func TestContainsDomain(t *testing.T) { tests := []struct { name string @@ -564,9 +752,198 @@ func TestContainsDomain(t *testing.T) { } } +func TestEnsureSliceContainsAll(t *testing.T) { + tests := []struct { + name string + data []string + mandatory []string + expected []string + }{ + { + name: "All elements present", + data: []string{"a", "b", "c"}, + mandatory: []string{"a", "b"}, + expected: []string{"a", "b", "c"}, + }, + { + name: "Some elements missing", + data: []string{"a", "b"}, + mandatory: []string{"a", "b", "c"}, + expected: []string{"a", "b", "c"}, + }, + { + name: "No elements present", + data: []string{}, + mandatory: []string{"a", "b", "c"}, + expected: []string{"a", "b", "c"}, + }, + { + name: "Empty mandatory list", + data: []string{"a", "b", "c"}, + mandatory: []string{}, + expected: []string{"a", "b", "c"}, + }, + { + name: "No elements in both lists", + data: []string{}, + mandatory: []string{}, + expected: []string{}, + }, + { + name: "Duplicate elements in mandatory list", + data: []string{"a", "b"}, + mandatory: []string{"b", "c", "c"}, + expected: []string{"a", "b", "c", "c"}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result := EnsureSliceContainsAll(test.data, test.mandatory) + require.ElementsMatch(t, test.expected, result, "resulting list should match expected value") + }) + } +} + +func TestSortUInt32s(t *testing.T) { + tests := []struct { + name string + data []uint32 + expected []uint32 + }{ + { + name: "Already sorted", + data: []uint32{1, 2, 3, 4, 5}, + expected: []uint32{1, 2, 3, 4, 5}, + }, + { + name: "Reverse order", + data: []uint32{5, 4, 3, 2, 1}, + expected: []uint32{1, 2, 3, 4, 5}, + }, + { + name: "Unsorted", + data: []uint32{3, 1, 4, 5, 2}, + expected: []uint32{1, 2, 3, 4, 5}, + }, + { + name: "With duplicates", + data: []uint32{3, 1, 4, 1, 5, 2, 3}, + expected: []uint32{1, 1, 2, 3, 3, 4, 5}, + }, + { + name: "Single element", + data: []uint32{1}, + expected: []uint32{1}, + }, + { + name: "Empty slice", + data: []uint32{}, + expected: []uint32{}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + SortUInt32s(test.data) + require.Equal(t, test.expected, test.data, "the sorted data should match the expected value") + }) + } +} + +func TestUInt32sAreSorted(t *testing.T) { + tests := []struct { + name string + data []uint32 + expected bool + }{ + { + name: "Sorted data", + data: []uint32{1, 2, 3, 4, 5}, + expected: true, + }, + { + name: "Unsorted data", + data: []uint32{5, 3, 4, 1, 2}, + expected: false, + }, + { + name: "Empty data", + data: []uint32{}, + expected: true, + }, + { + name: "Single element", + data: []uint32{42}, + expected: true, + }, + { + name: "All elements equal", + data: []uint32{7, 7, 7, 7}, + expected: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result := UInt32sAreSorted(test.data) + require.Equal(t, test.expected, result, "the result should match the expected value") + }) + } +} + +func TestValidateTimestamp(t *testing.T) { + tests := []struct { + name string + timestamp time.Time + expectedTime time.Time + replaced bool + }{ + { + name: "Valid timestamp", + timestamp: time.Date(2024, time.June, 3, 23, 24, 10, 0, time.Local), + expectedTime: time.Date(2024, time.June, 3, 23, 24, 10, 0, time.Local), + replaced: false, + }, + { + name: "Log Floating-Pont Timestamp", + timestamp: time.Unix(1517336108, int64((0.231879)*1e9)), // 1517336108.231879 + expectedTime: time.Unix(1517336108, 231879000), + replaced: false, + }, + { + name: "Unset Timestamp", + timestamp: time.Time{}, + expectedTime: time.Unix(0, 1), + replaced: true, + }, + { + name: "MaxInt64 timestamp", + timestamp: time.Unix(math.MaxInt64, 0), + expectedTime: time.Unix(0, 1), + replaced: true, + }, + { + name: "Negative timestamp", + timestamp: time.Unix(-1, 0), + expectedTime: time.Unix(0, 1), + replaced: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ts, replaced := ValidateTimestamp(test.timestamp) + require.Equal(t, test.expectedTime, ts, "timestamp should match expected value") + require.Equal(t, test.replaced, replaced, "replaced should match expected value") + }) + } +} + func TestParseRelativePath(t *testing.T) { home, err := os.UserHomeDir() require.NoError(t, err) + fmt.Println("home: ", home) workingDir, err := os.Getwd() require.NoError(t, err) @@ -577,39 +954,51 @@ func TestParseRelativePath(t *testing.T) { name string path string expected string - expectErr bool + expectErr error }{ { name: "Home directory", path: "~/data", expected: home + "/data", - expectErr: false, + expectErr: nil, + }, + { + name: "Current directory path", + path: "./", + expected: workingDir, + // expectedPath: filepath.Join(currentDir, "./mydir"), + expectErr: nil, }, { name: "Relative directory - 1 deep", path: "./data", expected: workingDir + "/data", - expectErr: false, + expectErr: nil, }, { name: "Relative directory - 2 deep", path: "../data", expected: currentDir + "/data", - expectErr: false, + expectErr: nil, }, { name: "Absolute path", path: "/home/logs", expected: "/home/logs", - expectErr: false, + expectErr: nil, + }, + { + name: "Empty path", + expected: "", + expectErr: ErrInvalidPath, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { result, err := ParseRelativePath(test.path) - if test.expectErr { - require.Error(t, err, "error was expected") + if test.expectErr != nil { + require.EqualError(t, err, test.expectErr.Error(), "error should match expected value") } else { require.NoError(t, err, "parsing relative path should not produce an error") require.Equal(t, test.expected, result, "relative path should match expected value, got: %s, expected: %s", result, test.expected) @@ -618,7 +1007,208 @@ func TestParseRelativePath(t *testing.T) { } } -// TestCheckForNewerVersion tests the CheckForNewerVersion function +func TestValidateDirectory(t *testing.T) { + tests := []struct { + name string + setup func(afs afero.Fs) + dir string + expectedError error + }{ + { + name: "Directory is Valid", + setup: func(afs afero.Fs) { + require.NoError(t, afs.Mkdir("/nonemptydir", 0755)) + require.NoError(t, afero.WriteFile(afs, "/nonemptydir/file.txt", []byte("content"), 0644)) + }, + dir: "/nonemptydir", + expectedError: nil, + }, + { + name: "Directory Does Not Exist", + setup: func(_ afero.Fs) {}, + dir: "/nonexistent", + expectedError: ErrDirDoesNotExist, + }, + { + name: "Path is Not a Directory", + setup: func(afs afero.Fs) { + require.NoError(t, afero.WriteFile(afs, "/file.txt", []byte("content"), 0644)) + }, + dir: "/file.txt", + expectedError: ErrPathIsNotDir, + }, + { + name: "Directory is Empty", + setup: func(afs afero.Fs) { + require.NoError(t, afs.Mkdir("/emptydir", 0755)) + }, + dir: "/emptydir", + expectedError: ErrDirIsEmpty, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + afs := afero.NewMemMapFs() + test.setup(afs) + + err := ValidateDirectory(afs, test.dir) + if test.expectedError != nil { + require.Error(t, err, "error should not be nil") + require.ErrorContains(t, err, test.expectedError.Error(), "error should contain expected value") + } else { + require.NoError(t, err, "validating directory should not produce an error") + } + }) + } +} + +func TestValidateFile(t *testing.T) { + tests := []struct { + name string + setup func(afs afero.Fs) + file string + expectedError error + }{ + { + name: "File is Valid", + setup: func(afs afero.Fs) { + require.NoError(t, afero.WriteFile(afs, "/file.txt", []byte("content"), 0644)) + }, + file: "/file.txt", + expectedError: nil, + }, + { + name: "File is Empty", + setup: func(afs afero.Fs) { + require.NoError(t, afero.WriteFile(afs, "/emptyfile.txt", []byte(""), 0644)) + }, + file: "/emptyfile.txt", + expectedError: ErrFileIsEmtpy, + }, + { + name: "File Does Not Exist", + setup: func(_ afero.Fs) {}, + file: "/nonexistent", + expectedError: ErrFileDoesNotExist, + }, + { + name: "Path is a Directory", + setup: func(afs afero.Fs) { + require.NoError(t, afs.Mkdir("/directory", 0755)) + }, + file: "/directory", + expectedError: ErrPathIsDir, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + afs := afero.NewMemMapFs() + test.setup(afs) + + err := ValidateFile(afs, test.file) + if test.expectedError != nil { + require.Error(t, err, "error should not be nil") + require.ErrorContains(t, err, test.expectedError.Error(), "error should contain expected value") + } else { + require.NoError(t, err, "validating file should not produce an error") + } + }) + } +} + +func TestValidatePath(t *testing.T) { + tests := []struct { + name string + setup func(afs afero.Fs) + path string + expected [3]bool // exists, isDir, isEmpty + expectedError error + }{ + { + name: "Path is Valid Non-Empty File", + setup: func(afs afero.Fs) { + require.NoError(t, afero.WriteFile(afs, "/file.txt", []byte("content"), 0644)) + }, + path: "/file.txt", + expected: [3]bool{true, false, false}, + expectedError: nil, + }, + { + name: "Path is Valid Empty File", + setup: func(afs afero.Fs) { + require.NoError(t, afero.WriteFile(afs, "/file.txt", []byte(""), 0644)) + }, + path: "/file.txt", + expected: [3]bool{true, false, true}, + expectedError: nil, + }, + { + name: "Path is Valid Non-Empty Directory", + setup: func(afs afero.Fs) { + require.NoError(t, afs.Mkdir("/nonemptydir", 0755)) + require.NoError(t, afero.WriteFile(afs, "/nonemptydir/file.txt", []byte("content"), 0644)) + }, + path: "/nonemptydir", + expected: [3]bool{true, true, false}, + expectedError: nil, + }, + { + name: "Path is Valid Empty Directory", + setup: func(afs afero.Fs) { + require.NoError(t, afs.Mkdir("/emptydir", 0755)) + }, + path: "/emptydir", + expected: [3]bool{true, true, true}, + expectedError: nil, + }, + { + name: "Non-Existent Path", + setup: func(_ afero.Fs) {}, + path: "/nonexistent", + expected: [3]bool{false, false, false}, + expectedError: nil, + }, + { + name: "Empty Path", + setup: func(_ afero.Fs) {}, + path: "", + expected: [3]bool{false, false, false}, + expectedError: ErrInvalidPath, + }, + { + name: "Nil filesystem", + setup: func(_ afero.Fs) {}, + path: "/some/path", + expected: [3]bool{false, false, false}, + expectedError: fmt.Errorf("filesystem is nil"), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var afs afero.Fs + if test.name != "Nil filesystem" { + afs = afero.NewMemMapFs() + } + test.setup(afs) + + exists, isDir, isEmpty, err := validatePath(afs, test.path) + + if test.expectedError != nil { + // require.Error(t, err) + require.ErrorContains(t, err, test.expectedError.Error(), "error should contain expected value") + } else { + require.NoError(t, err, "validating path should not produce an error") + require.Equal(t, test.expected[0], exists, "exist flag should be %v", test.expected[0]) + require.Equal(t, test.expected[1], isDir, "isDir flag should be %v", test.expected[1]) + require.Equal(t, test.expected[2], isEmpty, "isEmpty flag should be %v", test.expected[2]) + } + }) + } +} + func TestCheckForNewerVersion(t *testing.T) { tests := []struct { name string @@ -687,3 +1277,61 @@ func TestCheckForNewerVersion(t *testing.T) { }) } } + +func TestGetLatestReleaseVersion(t *testing.T) { + tests := []struct { + name string + owner string + repo string + latestVersion string + expected string + expectedError bool + }{ + { + name: "Valid Latest Release", + owner: "activecm", + repo: "rita", + latestVersion: "v2.0.0", + expected: "v2.0.0", + expectedError: false, + }, + { + name: "Error Fetching Latest Release", + owner: "activecm", + repo: "rita", + expected: "", + expectedError: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + // Create a test server + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if test.expectedError { + http.Error(w, "error", http.StatusInternalServerError) + } else { + fmt.Fprintf(w, `{"tag_name": "%s"}`, test.latestVersion) + } + })) + defer ts.Close() + + // Override the GitHub client base URL + client := github.NewClient(nil) + newBaseURL, err := client.BaseURL.Parse(ts.URL + "/") + require.NoError(t, err, "failed to parse base URL") + client.BaseURL = newBaseURL + + result, err := GetLatestReleaseVersion(client, test.owner, test.repo) + + if test.expectedError { + require.Error(t, err, "error should not be nil") + require.ErrorContains(t, err, "error fetching latest release", "error should contain expected value") + } else { + require.NoError(t, err, "fetching latest release should not produce an error") + require.Equal(t, test.expected, result, "the result should match the expected value") + } + + }) + } +} diff --git a/viewer/sidebar.go b/viewer/sidebar.go index 97f6dfe..90de634 100644 --- a/viewer/sidebar.go +++ b/viewer/sidebar.go @@ -42,25 +42,36 @@ func (m *sidebarModel) Init() tea.Cmd { m.Viewport.SetContent(m.getSidebarContents()) return nil } + +type UpdateItem *Item + func (m *sidebarModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { var cmds []tea.Cmd - switch msg.(type) { - case tea.KeyMsg: - // if k := msg.String(); k == "ctrl+c" || k == "q" || k == "esc" { - // return m, tea.Quit - // } + switch msg := msg.(type) { + + case *Item: + + m.Data = msg + content := m.getSidebarContents() + numlines := strings.Count(content, "\n") + 1 + 2 + + numToClear := m.Viewport.Height - numlines + if numToClear > 0 { + spaces := m.Viewport.Width - 2 + for i := 0; i < numToClear; i++ { + content += fmt.Sprintf("%*s\n", spaces, "") + } + } + + m.Viewport.SetContent(content) case tea.WindowSizeMsg: cmds = append(cmds, viewport.Sync(m.Viewport)) - - // return m, nil } return m, tea.Batch(cmds...) } func (m *sidebarModel) View() string { - m.Viewport.SetContent(m.getSidebarContents()) - m.Viewport.Height = m.Height borderColor := mauve if m.ScrollEnabled { borderColor = green @@ -71,13 +82,17 @@ func (m *sidebarModel) View() string { Padding(0, 1). Border(lipgloss.RoundedBorder()). BorderForeground(borderColor) - hi := style.Render(m.Viewport.View()) - return lipgloss.NewStyle().Render(hi) - // Border(lipgloss.NormalBorder()) + sidebar := style.Render(m.Viewport.View()) + return lipgloss.NewStyle().Render(sidebar) + } // getSidebarContents gets and formats the data to be displayed in the sidebar func (m *sidebarModel) getSidebarContents() string { + if m.Data == nil { + return lipgloss.NewStyle().Foreground(overlay0).Render("No result found.") + } + // get header var target string headerPadding := 2 @@ -119,15 +134,17 @@ func (m *sidebarModel) getSidebarContents() string { connInfoLabel := sectionStyle.Render("「 Connection Info 」") + dataStyle := lipgloss.NewStyle().Foreground(defaultTextColor) + // get connection count connCountStyle := lipgloss.NewStyle().Background(overlay2).Foreground(base).Bold(true).Padding(0, 2) connCountHeader := connCountStyle.Render("Connection Count") - connCount := lipgloss.JoinVertical(lipgloss.Top, connCountHeader, fmt.Sprintf("%d", m.Data.Count)) + connCount := dataStyle.Render(lipgloss.JoinVertical(lipgloss.Top, connCountHeader, fmt.Sprintf("%d", m.Data.Count))) // get total bytes bytesHeaderStyle := lipgloss.NewStyle().Background(overlay2).Foreground(base).Bold(true).Padding(0, 2) bytesHeader := bytesHeaderStyle.Render("Total Bytes") - bytes := lipgloss.JoinVertical(lipgloss.Top, bytesHeader, m.Data.TotalBytesFormatted) + bytes := dataStyle.Render(lipgloss.JoinVertical(lipgloss.Top, bytesHeader, m.Data.TotalBytesFormatted)) // get port:proto:service portProtoService := m.Data.GetPortProtoService() @@ -145,7 +162,7 @@ func (m *sidebarModel) getSidebarContents() string { // render header portsHeader := portsHeaderStyle.Render("Port : Proto : Service") - ports = lipgloss.JoinVertical(lipgloss.Top, portsHeader, strings.Join(portProtoService, ",")) + ports = dataStyle.Render(lipgloss.JoinVertical(lipgloss.Top, portsHeader, strings.Join(portProtoService, "\n"))) // strings.Join(portProtoService, "\n") // calculate the number of lines available for port data // remainingLines := m.viewport.Height - (lipgloss.Height(heading) + lipgloss.Height(modifiers) + lipgloss.Height(modifierLabel) + lipgloss.Height(connInfoLabel) + lipgloss.Height(bytes) + lipgloss.Height(connCount)) @@ -182,7 +199,7 @@ func (m *sidebarModel) renderModifiers() string { width := lipgloss.Width(newMod) if m.Viewport.Width <= width { - modifierLines = append(modifierLines, lipgloss.JoinHorizontal(lipgloss.Left, linebreakStyle.Render(currentModifiers))) + modifierLines = append(modifierLines, lipgloss.NewStyle().Foreground(defaultTextColor).Render(lipgloss.JoinHorizontal(lipgloss.Left, linebreakStyle.Render(currentModifiers)))) currentModifiers = mod if i != len(renderedModifiers)-1 { currentModifiers = newlineStyle.Render(mod) diff --git a/viewer/viewer.go b/viewer/viewer.go index 2205cb5..8353e19 100644 --- a/viewer/viewer.go +++ b/viewer/viewer.go @@ -189,7 +189,7 @@ func (m *Model) Init() tea.Cmd { // Update updates the model func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { - var cmd tea.Cmd + var cmds []tea.Cmd switch msg := msg.(type) { case tea.WindowSizeMsg: // make the footer the entire width of the terminal @@ -231,7 +231,8 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { // handle filtering case m.SearchBar.TextInput.Focused(): - cmd = m.handleFiltering(msg) + cmd := m.handleFiltering(msg) + cmds = append(cmds, cmd) // clear filtering (when search bar not focused) case key.Matches(msg, m.keys.clearFilter): @@ -239,18 +240,24 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { // handle quiting case key.Matches(msg, m.keys.quit): - cmd = tea.Quit + cmd := tea.Quit + cmds = append(cmds, cmd) // otherwise, handle browsing default: - cmd = m.handleBrowsing(msg) + cmd := m.handleBrowsing(msg) + cmds = append(cmds, cmd) } case StillLoadingResults, FooterFlash: - _, cmd = m.Footer.Update(msg) + _, cmd := m.Footer.Update(msg) + cmds = append(cmds, cmd) case FinishedLoadingResults: case spinner.TickMsg: + var cmd tea.Cmd m.Footer.spinner, cmd = m.Footer.spinner.Update(msg) + cmds = append(cmds, cmd) + } // update sidebar @@ -271,15 +278,17 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { // set sidebar data to the selected item if data, ok := m.List.Rows.Items()[m.List.Rows.Index()].(*Item); ok { - m.SideBar.Data = data + _, cmd := m.SideBar.Update(data) + cmds = append(cmds, cmd) } } else { - // if there are no items to display, set the sidebar data to an empty item - m.SideBar.Data = &Item{} + // if there are no items to display, set the sidebar data to nil + _, cmd := m.SideBar.Update(nil) + cmds = append(cmds, cmd) } - return m, cmd + return m, tea.Batch(cmds...) } // View renders the model to the terminal @@ -429,6 +438,11 @@ func (m *Model) requestResults(appendResults bool) { m.Footer.loading = true // time.Sleep(4 * time.Second) + // reset the server page number if we're not appending results to the rows list + if !appendResults { + m.serverPage = 0 + } + // get results from database items, appliedFilter, err := GetResults(m.db, filter, m.serverPage, m.serverPageSize, m.minTS) if err != nil {