mirror of
https://github.com/activecm/rita
synced 2026-06-08 13:02:45 +00:00
Liza Tsibur
664 lines
29 KiB
Go
664 lines
29 KiB
Go
package database_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/activecm/rita/v5/cmd"
|
|
"github.com/activecm/rita/v5/database"
|
|
"github.com/activecm/rita/v5/util"
|
|
|
|
clickhouse "github.com/ClickHouse/clickhouse-go/v2"
|
|
"github.com/spf13/afero"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func (d *DatabaseTestSuite) TestConnectToDB() {
|
|
|
|
// create and connect to a new database
|
|
d.Run("Connect to Existing Database", func() {
|
|
t := d.T()
|
|
_, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB", false, false)
|
|
require.NoError(t, err, "importing data should not produce an error")
|
|
|
|
db, err := database.ConnectToDB(context.Background(), "testDB", d.cfg, nil)
|
|
require.NoError(t, err, "connecting to createddatabase should not produce an error")
|
|
require.NotNil(t, db)
|
|
})
|
|
|
|
// attempt to connect to a non-existent database
|
|
d.Run("Connect to Non-Existent Database", func() {
|
|
t := d.T()
|
|
db, err := database.ConnectToDB(context.Background(), "nonExistentDB", d.cfg, nil)
|
|
require.Error(t, err, "connecting to a non-existent database should produce an error")
|
|
require.Nil(t, db)
|
|
})
|
|
|
|
// attempt to connect with invalid configuration
|
|
d.Run("Invalid Configuration", func() {
|
|
t := d.T()
|
|
invalidCfg := *d.cfg
|
|
invalidCfg.Env.DBConnection = "invalid connection string"
|
|
|
|
db, err := database.ConnectToDB(context.Background(), "testDB", &invalidCfg, nil)
|
|
require.Error(t, err, "connecting with invalid configuration should produce an error")
|
|
require.Nil(t, db)
|
|
})
|
|
|
|
// attempt to connect with a cancelled context
|
|
d.Run("Cancel Context", func() {
|
|
t := d.T()
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
// attempt to connect with a cancelled context
|
|
db, err := database.ConnectToDB(ctx, "testDB", d.cfg, nil)
|
|
require.Error(t, err, "connecting with a cancelled context should produce an error")
|
|
require.Nil(t, db)
|
|
})
|
|
}
|
|
|
|
func (d *DatabaseTestSuite) TestMinMaxTimestamps() {
|
|
d.Run("24 Hour Dataset", func() {
|
|
t := d.T()
|
|
// import a dataset with 24 hours of data
|
|
results, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB", false, false)
|
|
require.NoError(t, err, "importing data should not produce an error")
|
|
|
|
// connect to the database
|
|
db, err := database.ConnectToDB(context.Background(), "testDB", d.cfg, nil)
|
|
require.NoError(t, err, "connecting to database should not produce an error")
|
|
|
|
// get the min and max timestamps from a test function that queries the conn, openconn, and dns tables
|
|
_, max := getMinMaxTimestampsFromTables(t, db, results.ImportID[0])
|
|
require.NoError(t, err, "getting min/max timestamps should not error")
|
|
|
|
// fetch the beacon minimum and maximum timestamps
|
|
minTSBeacon, maxTSBeacon, _, err := db.GetBeaconMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching beacon min and max timestamps should not produce an error")
|
|
|
|
// fetch the true minimum and maximum timestamps
|
|
minTS, maxTS, _, useCurrentTime, err := db.GetTrueMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching min and max timestamps should not produce an error")
|
|
|
|
// validate that the current time is not used
|
|
require.False(t, useCurrentTime, "useCurrentTime should be false")
|
|
|
|
// check that the timestamps are within the expected range
|
|
require.True(t, maxTS.After(minTS), "max timestamp should be after min timestamp")
|
|
|
|
// capped min timestamp
|
|
minTSCapped := max.Add(-24 * time.Hour)
|
|
require.InDelta(t, 24.0, maxTS.Sub(minTS).Hours(), 0.1, "timestamp difference should be close to 24 hours")
|
|
|
|
fmt.Println("Max Timestamp: ", maxTS)
|
|
fmt.Println("Min Timestamp: ", minTS)
|
|
|
|
// check that the true timestamps match the min and max timestamps from the test function
|
|
require.Equal(t, minTSCapped, minTS, "min timestamp should match min timestamp from the test function")
|
|
require.Equal(t, max, maxTS, "max timestamp should match max timestamp from the test function")
|
|
|
|
// check that the beacon timestamps are within the expected range
|
|
require.True(t, maxTSBeacon.Before(max) || maxTSBeacon.Equal(max), "max beacon timestamp should be before or equal to max timestamp from the test function")
|
|
require.True(t, minTSBeacon.After(minTSCapped) || minTSBeacon.Equal(minTSCapped), "min beacon timestamp should be after or equal to min timestamp from the test function")
|
|
})
|
|
|
|
d.Run("Less Than 24 Hours of Data", func() {
|
|
t := d.T()
|
|
// import a dataset with < 24 hours of data
|
|
results, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/proxy_rolling", "testDB", false, false)
|
|
require.NoError(t, err, "importing data should not produce an error")
|
|
|
|
// connect to the database
|
|
db, err := database.ConnectToDB(context.Background(), "testDB", d.cfg, nil)
|
|
require.NoError(t, err, "connecting to database should not produce an error")
|
|
|
|
// get the min and max timestamps from a test function that queries the conn, openconn, and dns tables
|
|
min, max := getMinMaxTimestampsFromTables(t, db, results.ImportID[0])
|
|
require.NoError(t, err, "getting min/max timestamps should not error")
|
|
|
|
// fetch the beacon minimum and maximum timestamps
|
|
minTSBeacon, maxTSBeacon, _, err := db.GetBeaconMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching beacon min and max timestamps should not produce an error")
|
|
|
|
// fetch the true minimum and maximum timestamps
|
|
minTS, maxTS, _, useCurrentTime, err := db.GetTrueMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching min and max timestamps should not produce an error")
|
|
|
|
// validate that the current time is not used
|
|
require.False(t, useCurrentTime, "useCurrentTime should be false")
|
|
|
|
// check that the timestamps are within the expected range
|
|
require.True(t, maxTS.After(minTS), "max timestamp should be after min timestamp")
|
|
require.Less(t, maxTS.Sub(minTS).Hours(), 24.0, "timestamp difference should be less than 24 hours")
|
|
|
|
// check that the max timestamp matches the max timestamps from the test function
|
|
require.Equal(t, max, maxTS, "max timestamp should match max timestamp from test function")
|
|
|
|
// since the dataset is less than 24 hours, the min timestamp should match the min timestamp from the test function
|
|
require.Equal(t, min, minTS, "min timestamp should match min timestamp from test function")
|
|
|
|
// check that the beacon timestamps are within the expected range
|
|
require.True(t, maxTSBeacon.Before(max) || maxTSBeacon.Equal(max), "max beacon timestamp should be before or equal to max timestamp from the test function")
|
|
require.True(t, minTSBeacon.After(min) || minTSBeacon.Equal(min), "min beacon timestamp should be after or equal to min timestamp from the test function")
|
|
})
|
|
|
|
d.Run("Greater Than 24 Hours of Data", func() {
|
|
t := d.T()
|
|
|
|
afs := afero.NewMemMapFs()
|
|
|
|
// create directory
|
|
directory := "/logs"
|
|
err := afs.Mkdir(directory, os.FileMode(0o775))
|
|
require.NoError(t, err, "creating directory should not produce an error")
|
|
|
|
// create mock data file
|
|
fileName := "conn.log"
|
|
path := filepath.Join(directory, fileName)
|
|
file, err := afs.Create(path)
|
|
require.NoError(t, err, "creating file should not produce an error")
|
|
|
|
// set file permissions
|
|
err = afs.Chmod(path, os.FileMode(0o775))
|
|
require.NoError(t, err, "changing file permissions should not produce an error")
|
|
|
|
// generate timestamp from current time and format to the timestamp format used in the log file
|
|
currentTime := time.Now().UTC()
|
|
formattedMaxTime := fmt.Sprintf("%d.%06d", currentTime.Unix(), currentTime.Nanosecond()/1000)
|
|
|
|
// create interval timestamps of 4 hours ago
|
|
fourHoursAgo := currentTime.Add(-4 * time.Hour)
|
|
formattedTime := fmt.Sprintf("%d.%06d", fourHoursAgo.Unix(), fourHoursAgo.Nanosecond()/1000)
|
|
|
|
// create timestamp for 48 hours ago
|
|
twoDaysAgo := currentTime.Add(-48 * time.Hour)
|
|
formattedMinTime := fmt.Sprintf("%d.%06d", twoDaysAgo.Unix(), twoDaysAgo.Nanosecond()/1000)
|
|
|
|
// create mock data
|
|
log := []byte("#separator \\x09\n" +
|
|
"#set_separator\t,\n" +
|
|
"#empty_field\t(empty)\n" +
|
|
"#unset_field\t-\n" +
|
|
"#path\tconn\n" +
|
|
"#open\t2019-02-28-12-07-01\n" +
|
|
"#fields\tts\tuid\tid.orig_h\tid.resp_h\n" +
|
|
"#types\ttime\tstring\taddr\taddr\n" +
|
|
formattedMaxTime + "\tCxT121\t10.0.0.1\t52.12.0.1\n" +
|
|
formattedTime + "\tCxT121\t10.0.0.1\t52.12.0.1\n" +
|
|
formattedTime + "\tCxT122\t10.0.0.2\t52.12.0.2\n" +
|
|
formattedTime + "\tCxT122\t10.0.0.2\t52.12.0.2\n" +
|
|
formattedTime + "\tCxT123\t10.0.0.3\t52.12.0.3\n" +
|
|
formattedTime + "\tCxT123\t10.0.0.3\t52.12.0.3\n" +
|
|
formattedTime + "\tCxT124\t10.0.0.4\t52.12.0.4\n" +
|
|
formattedTime + "\tCxT124\t10.0.0.4\t52.12.0.4\n" +
|
|
formattedTime + "\tCxT125\t10.0.0.5\t52.12.0.5\n" +
|
|
formattedMinTime + "\tCxT125\t10.0.0.5\t52.12.0.5\n",
|
|
)
|
|
bytesWritten, err := file.Write(log)
|
|
require.NoError(t, err, "writing data to file should not produce an error")
|
|
require.Equal(t, len(log), bytesWritten, "number of bytes written should be equal to the length of the log data")
|
|
|
|
err = file.Close()
|
|
require.NoError(t, err, "closing file should not produce an error")
|
|
|
|
// import the mock data
|
|
results, err := cmd.RunImportCmd(time.Now(), d.cfg, afs, directory, "testDB", false, false)
|
|
require.NoError(t, err, "importing data should not produce an error")
|
|
|
|
// connect to the database
|
|
db, err := database.ConnectToDB(context.Background(), "testDB", d.cfg, nil)
|
|
require.NoError(t, err, "connecting to database should not produce an error")
|
|
|
|
// get the min and max timestamps from a test function that queries the conn, openconn, and dns tables
|
|
min, max := getMinMaxTimestampsFromTables(t, db, results.ImportID[0])
|
|
require.NoError(t, err, "getting min/max timestamps should not error")
|
|
|
|
// fetch the beacon minimum and maximum timestamps
|
|
minTSBeacon, maxTSBeacon, _, err := db.GetBeaconMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching beacon min and max timestamps should not produce an error")
|
|
|
|
// fetch the true minimum and maximum timestamps
|
|
minTS, maxTS, _, useCurrentTime, err := db.GetTrueMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching min and max timestamps should not produce an error")
|
|
|
|
// validate that the current time is not used
|
|
require.False(t, useCurrentTime, "useCurrentTime should be false")
|
|
|
|
// check that the timestamps are within the expected range
|
|
require.True(t, maxTS.After(minTS), "max timestamp should be after min timestamp")
|
|
require.InDelta(t, 24.0, maxTS.Sub(minTS).Hours(), 0.1, "timestamp difference should be close to 24 hours")
|
|
|
|
// check that the max timestamp matches the max timestamps from the test function
|
|
require.Equal(t, max, maxTS, "max timestamp should match max timestamp from test function")
|
|
|
|
// since the dataset is > 24 hours, the min timestamp will get capped to 24 hours from the max timestamp
|
|
require.NotEqual(t, min, minTS, "min timestamp should not match min timestamp from test function")
|
|
require.Equal(t, maxTS.Add(-24*time.Hour), minTS, "min timestamp should be 24 hours from max timestamp")
|
|
|
|
// check that the beacon timestamps are within the expected range
|
|
require.True(t, maxTSBeacon.Before(max) || maxTSBeacon.Equal(max), "max beacon timestamp should be before or equal to max timestamp from the test function")
|
|
require.True(t, minTSBeacon.After(maxTS.Add(-24*time.Hour)) || minTSBeacon.Equal(maxTS.Add(-24*time.Hour)), "min beacon timestamp should be after or equal to min timestamp from the test function")
|
|
})
|
|
|
|
d.Run("Rolling, Max Timestamp < 24Hrs Ago", func() {
|
|
t := d.T()
|
|
afs := afero.NewMemMapFs()
|
|
|
|
// create directory
|
|
directory := "/logs"
|
|
err := afs.Mkdir(directory, os.FileMode(0o775))
|
|
require.NoError(t, err, "creating directory should not produce an error")
|
|
|
|
// create mock data file
|
|
fileName := "conn.log"
|
|
path := filepath.Join(directory, fileName)
|
|
file, err := afs.Create(path)
|
|
require.NoError(t, err, "creating file should not produce an error")
|
|
|
|
// set file permissions
|
|
err = afs.Chmod(path, os.FileMode(0o775))
|
|
require.NoError(t, err, "changing file permissions should not produce an error")
|
|
|
|
// generate timestamp from current time and format to the timestamp format used in the log file
|
|
maxTime := time.Now().UTC().Add(-1 * time.Hour)
|
|
formattedMaxTime := fmt.Sprintf("%d.%06d", maxTime.Unix(), maxTime.Nanosecond()/1000)
|
|
|
|
// create interval timestamps of 4 hours ago
|
|
fourHoursAgo := maxTime.Add(-4 * time.Hour)
|
|
formattedTime := fmt.Sprintf("%d.%06d", fourHoursAgo.Unix(), fourHoursAgo.Nanosecond()/1000)
|
|
|
|
// create timestamp for 24 hours ago
|
|
minTime := maxTime.Add(-24 * time.Hour)
|
|
formattedMinTime := fmt.Sprintf("%d.%06d", minTime.Unix(), minTime.Nanosecond()/1000)
|
|
|
|
// create mock data
|
|
log := []byte("#separator \\x09\n" +
|
|
"#set_separator\t,\n" +
|
|
"#empty_field\t(empty)\n" +
|
|
"#unset_field\t-\n" +
|
|
"#path\tconn\n" +
|
|
"#open\t2019-02-28-12-07-01\n" +
|
|
"#fields\tts\tuid\tid.orig_h\tid.resp_h\n" +
|
|
"#types\ttime\tstring\taddr\taddr\n" +
|
|
formattedMaxTime + "\tCxT121\t10.0.0.1\t52.12.0.1\n" +
|
|
formattedTime + "\tCxT121\t10.0.0.1\t52.12.0.1\n" +
|
|
formattedTime + "\tCxT122\t10.0.0.2\t52.12.0.2\n" +
|
|
formattedTime + "\tCxT122\t10.0.0.2\t52.12.0.2\n" +
|
|
formattedTime + "\tCxT123\t10.0.0.3\t52.12.0.3\n" +
|
|
formattedTime + "\tCxT123\t10.0.0.3\t52.12.0.3\n" +
|
|
formattedTime + "\tCxT124\t10.0.0.4\t52.12.0.4\n" +
|
|
formattedTime + "\tCxT124\t10.0.0.4\t52.12.0.4\n" +
|
|
formattedTime + "\tCxT125\t10.0.0.5\t52.12.0.5\n" +
|
|
formattedMinTime + "\tCxT125\t10.0.0.5\t52.12.0.5\n",
|
|
)
|
|
bytesWritten, err := file.Write(log)
|
|
require.NoError(t, err, "writing data to file should not produce an error")
|
|
require.Equal(t, len(log), bytesWritten, "number of bytes written should be equal to the length of the log data")
|
|
|
|
err = file.Close()
|
|
require.NoError(t, err, "closing file should not produce an error")
|
|
|
|
// import the mock data
|
|
results, err := cmd.RunImportCmd(time.Now(), d.cfg, afs, directory, "testDB", true, false)
|
|
require.NoError(t, err, "importing data should not produce an error")
|
|
|
|
// connect to the database
|
|
db, err := database.ConnectToDB(context.Background(), "testDB", d.cfg, nil)
|
|
require.NoError(t, err, "connecting to database should not produce an error")
|
|
|
|
// get the min and max timestamps from a test function that queries the conn, openconn, and dns tables
|
|
min, max := getMinMaxTimestampsFromTables(t, db, results.ImportID[0])
|
|
require.NoError(t, err, "getting min/max timestamps should not error")
|
|
|
|
// fetch the beacon minimum and maximum timestamps
|
|
minTSBeacon, maxTSBeacon, _, err := db.GetBeaconMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching beacon min and max timestamps should not produce an error")
|
|
|
|
// fetch the true minimum and maximum timestamps
|
|
minTS, maxTS, _, useCurrentTime, err := db.GetTrueMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching min and max timestamps should not produce an error")
|
|
|
|
// validate that the current time should be used since the max timestamp is <= 24 hours ago and the dataset is rolling
|
|
require.True(t, useCurrentTime, "useCurrentTime should be true")
|
|
|
|
// check that the timestamps are within the expected range
|
|
require.True(t, maxTS.After(minTS), "max timestamp should be after min timestamp")
|
|
require.InDelta(t, 24.0, maxTS.Sub(minTS).Hours(), 0.1, "timestamp difference should be close to 24 hours")
|
|
|
|
// check that the timestamps match the min and max timestamps from the test function
|
|
require.Equal(t, min, minTS, "min timestamp should match min timestamp from the test function")
|
|
require.Equal(t, max, maxTS, "max timestamp should match max timestamp from the test function")
|
|
|
|
// check that the beacon timestamps are within the expected range
|
|
require.True(t, maxTSBeacon.Before(max) || maxTSBeacon.Equal(max), "max beacon timestamp should be before or equal to max timestamp from the test function")
|
|
require.True(t, minTSBeacon.After(min) || minTSBeacon.Equal(min), "min beacon timestamp should be after or equal to min timestamp from the test function")
|
|
})
|
|
|
|
d.Run("Rolling, Max Timestamp > 24Hrs Ago", func() {
|
|
t := d.T()
|
|
afs := afero.NewMemMapFs()
|
|
|
|
// create directory
|
|
directory := "/logs"
|
|
err := afs.Mkdir(directory, os.FileMode(0o775))
|
|
require.NoError(t, err, "creating directory should not produce an error")
|
|
|
|
// create mock data file
|
|
fileName := "conn.log"
|
|
path := filepath.Join(directory, fileName)
|
|
file, err := afs.Create(path)
|
|
require.NoError(t, err, "creating file should not produce an error")
|
|
|
|
// set file permissions
|
|
err = afs.Chmod(path, os.FileMode(0o775))
|
|
require.NoError(t, err, "changing file permissions should not produce an error")
|
|
|
|
// generate timestamp from current time and format to the timestamp format used in the log file
|
|
maxTime := time.Now().UTC().Add(-25 * time.Hour)
|
|
formattedMaxTime := fmt.Sprintf("%d.%06d", maxTime.Unix(), maxTime.Nanosecond()/1000)
|
|
|
|
// create interval timestamps of 4 hours ago
|
|
fourHoursAgo := maxTime.Add(-4 * time.Hour)
|
|
formattedTime := fmt.Sprintf("%d.%06d", fourHoursAgo.Unix(), fourHoursAgo.Nanosecond()/1000)
|
|
|
|
// create timestamp for 48 hours ago
|
|
minTime := maxTime.Add(-48 * time.Hour)
|
|
formattedMinTime := fmt.Sprintf("%d.%06d", minTime.Unix(), minTime.Nanosecond()/1000)
|
|
|
|
// create mock data
|
|
log := []byte("#separator \\x09\n" +
|
|
"#set_separator\t,\n" +
|
|
"#empty_field\t(empty)\n" +
|
|
"#unset_field\t-\n" +
|
|
"#path\tconn\n" +
|
|
"#open\t2019-02-28-12-07-01\n" +
|
|
"#fields\tts\tuid\tid.orig_h\tid.resp_h\n" +
|
|
"#types\ttime\tstring\taddr\taddr\n" +
|
|
formattedMaxTime + "\tCxT121\t10.0.0.1\t52.12.0.1\n" +
|
|
formattedTime + "\tCxT121\t10.0.0.1\t52.12.0.1\n" +
|
|
formattedTime + "\tCxT122\t10.0.0.2\t52.12.0.2\n" +
|
|
formattedTime + "\tCxT122\t10.0.0.2\t52.12.0.2\n" +
|
|
formattedTime + "\tCxT123\t10.0.0.3\t52.12.0.3\n" +
|
|
formattedTime + "\tCxT123\t10.0.0.3\t52.12.0.3\n" +
|
|
formattedTime + "\tCxT124\t10.0.0.4\t52.12.0.4\n" +
|
|
formattedTime + "\tCxT124\t10.0.0.4\t52.12.0.4\n" +
|
|
formattedTime + "\tCxT125\t10.0.0.5\t52.12.0.5\n" +
|
|
formattedMinTime + "\tCxT125\t10.0.0.5\t52.12.0.5\n",
|
|
)
|
|
bytesWritten, err := file.Write(log)
|
|
require.NoError(t, err, "writing data to file should not produce an error")
|
|
require.Equal(t, len(log), bytesWritten, "number of bytes written should be equal to the length of the log data")
|
|
|
|
err = file.Close()
|
|
require.NoError(t, err, "closing file should not produce an error")
|
|
|
|
// import the mock data
|
|
results, err := cmd.RunImportCmd(time.Now(), d.cfg, afs, directory, "testDB", true, false)
|
|
require.NoError(t, err, "importing data should not produce an error")
|
|
|
|
// connect to the database
|
|
db, err := database.ConnectToDB(context.Background(), "testDB", d.cfg, nil)
|
|
require.NoError(t, err, "connecting to database should not produce an error")
|
|
|
|
// get the min and max timestamps from a test function that queries the conn, openconn, and dns tables
|
|
min, max := getMinMaxTimestampsFromTables(t, db, results.ImportID[0])
|
|
require.NoError(t, err, "getting min/max timestamps should not error")
|
|
|
|
// fetch the beacon minimum and maximum timestamps
|
|
minTSBeacon, maxTSBeacon, _, err := db.GetBeaconMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching beacon min and max timestamps should not produce an error")
|
|
|
|
// fetch the true minimum and maximum timestamps
|
|
minTS, maxTS, _, useCurrentTime, err := db.GetTrueMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching min and max timestamps should not produce an error")
|
|
|
|
// validate that the current time is not used
|
|
require.False(t, useCurrentTime, "useCurrentTime should be false")
|
|
|
|
// check that the timestamps are within the expected range
|
|
require.True(t, maxTS.After(minTS), "max timestamp should be after min timestamp")
|
|
require.InDelta(t, 24.0, maxTS.Sub(minTS).Hours(), 0.1, "timestamp difference should be close to 24 hours")
|
|
|
|
// check that the max timestamp matches the max timestamps from the test function
|
|
require.Equal(t, max, maxTS, "max timestamp should match max timestamp from test function")
|
|
|
|
// since the dataset is > 24 hours, the min timestamp will get capped to 24 hours from the max timestamp
|
|
require.NotEqual(t, min, minTS, "min timestamp should not match min timestamp from test function")
|
|
require.Equal(t, maxTS.Add(-24*time.Hour), minTS, "min timestamp should be 24 hours from max timestamp")
|
|
|
|
// check that the beacon timestamps are within the expected range
|
|
require.True(t, maxTSBeacon.Before(max) || maxTSBeacon.Equal(max), "max beacon timestamp should be before or equal to max timestamp from the test function")
|
|
require.True(t, minTSBeacon.After(maxTS.Add(-24*time.Hour)) || minTSBeacon.Equal(maxTS.Add(-24*time.Hour)), "min beacon timestamp should be after or equal to min timestamp from the test function")
|
|
})
|
|
|
|
d.Run("Non-Existent Database / Invalid Connection", func() {
|
|
t := d.T()
|
|
db := database.DB{}
|
|
|
|
// fetch the beacon minimum and maximum timestamps
|
|
minTSBeacon, maxTSBeacon, _, err := db.GetBeaconMinMaxTimestamps()
|
|
require.Error(t, err, "fetching beacon min and max timestamps should produce an error")
|
|
require.Equal(t, minTSBeacon, time.Unix(0, 0), "max timestamp should be zero unix time")
|
|
require.Equal(t, maxTSBeacon, time.Unix(0, 0), "min timestamp should be zero unix time")
|
|
require.Equal(t, err, database.ErrInvalidDatabaseConnection, "error should be invalid database connection")
|
|
|
|
// fetch the true minimum and maximum timestamps
|
|
minTS, maxTS, _, useCurrentTime, err := db.GetTrueMinMaxTimestamps()
|
|
require.Error(t, err, "fetching min and max timestamps should produce an error")
|
|
|
|
require.Equal(t, maxTS, time.Unix(0, 0), "max timestamp should be zero unix time")
|
|
require.Equal(t, minTS, time.Unix(0, 0), "min timestamp should be zero unix time")
|
|
require.False(t, useCurrentTime, "value passed back to useCurrentTime should be false")
|
|
require.Equal(t, err, database.ErrInvalidDatabaseConnection, "error should be invalid database connection")
|
|
|
|
})
|
|
}
|
|
|
|
func getMinMaxTimestampsFromTables(t *testing.T, db *database.DB, importID util.FixedString) (time.Time, time.Time) { // uses valid dataset
|
|
t.Helper()
|
|
|
|
type minMaxRes struct {
|
|
Min time.Time `ch:"min_timestamp"`
|
|
Max time.Time `ch:"max_timestamp"`
|
|
}
|
|
|
|
var result minMaxRes
|
|
|
|
// set context with importID and database parameters
|
|
ctx := clickhouse.Context(db.GetContext(), clickhouse.WithParameters(clickhouse.Parameters{
|
|
"importID": importID.Hex(),
|
|
"database": db.GetSelectedDB(),
|
|
}))
|
|
|
|
// get min and max timestamps from the conn, openconn, and udns tables
|
|
err := db.Conn.QueryRow(ctx, `
|
|
SELECT MIN(ts) AS min_timestamp, MAX(ts) AS max_timestamp
|
|
FROM (
|
|
SELECT ts FROM {database:Identifier}.conn
|
|
UNION ALL
|
|
SELECT ts FROM {database:Identifier}.openconn
|
|
UNION ALL
|
|
SELECT hour as ts FROM {database:Identifier}.udns
|
|
) AS combined_timestamps
|
|
`).ScanStruct(&result)
|
|
require.NoError(t, err)
|
|
|
|
fmt.Println("Min Timestamp: ", result.Min)
|
|
fmt.Println("Max Timestamp: ", result.Max)
|
|
|
|
return result.Min, result.Max
|
|
}
|
|
|
|
func (d *DatabaseTestSuite) TestGetNetworkSize() {
|
|
d.Run("Valid TSV Dataset", func() {
|
|
t := d.T()
|
|
// import a dataset
|
|
_, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB", false, false)
|
|
require.NoError(t, err, "importing data should not produce an error")
|
|
|
|
// connect to the database
|
|
db, err := database.ConnectToDB(context.Background(), "testDB", d.cfg, nil)
|
|
require.NoError(t, err, "connecting to database should not produce an error")
|
|
|
|
var result struct {
|
|
Count uint64 `ch:"count"`
|
|
}
|
|
|
|
// verify http linking wrote no more than 20 records with the same zeek uid
|
|
err = db.Conn.QueryRow(db.GetContext(), `
|
|
SELECT count() as count FROM conn
|
|
`).ScanStruct(&result)
|
|
require.NoError(t, err, "querying conn table should not produce an error")
|
|
|
|
fmt.Println("Count: ", result.Count)
|
|
|
|
// get the min timestamp
|
|
minTS, _, notFromConn, err := db.GetBeaconMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching min timestamp should not produce an error")
|
|
|
|
// validate which table min max is from
|
|
require.False(t, notFromConn, "min and max timestamps should be from conn table")
|
|
|
|
// get the network size (number of unique internal hosts for the past 24 hours)
|
|
networkSize, err := db.GetNetworkSize(minTS)
|
|
require.NoError(t, err, "getting network size should not produce an error")
|
|
|
|
// verify the expected network size
|
|
require.Greater(t, networkSize, uint64(0), "network size should be greater than zero")
|
|
|
|
require.Equal(t, 15, int(networkSize), "network size for valid tsv dataset should match expected value")
|
|
})
|
|
|
|
d.Run("Mocked Log File with Duplicate Internal Hosts", func() {
|
|
t := d.T()
|
|
afs := afero.NewMemMapFs()
|
|
|
|
// create directory
|
|
directory := "/logs"
|
|
err := afs.Mkdir(directory, os.FileMode(0o775))
|
|
require.NoError(t, err, "creating directory should not produce an error")
|
|
|
|
// create mock data file
|
|
fileName := "conn.log"
|
|
path := filepath.Join(directory, fileName)
|
|
file, err := afs.Create(path)
|
|
require.NoError(t, err, "creating file should not produce an error")
|
|
|
|
// set file permissions
|
|
err = afs.Chmod(path, os.FileMode(0o775))
|
|
require.NoError(t, err, "changing file permissions should not produce an error")
|
|
|
|
// create mock data with duplicates and write data to file
|
|
log := []byte("#separator \\x09\n" +
|
|
"#set_separator\t,\n" +
|
|
"#empty_field\t(empty)\n" +
|
|
"#unset_field\t-\n" +
|
|
"#path\tconn\n" +
|
|
"#open\t2019-02-28-12-07-01\n" +
|
|
"#fields\tts\tuid\tid.orig_h\tid.resp_h\n" +
|
|
"#types\ttime\tstring\taddr\taddr\n" +
|
|
"1715640994.367201\tCxT121\t10.0.0.1\t52.12.0.1\n" +
|
|
"1715640994.367201\tCxT121\t10.0.0.1\t52.12.0.1\n" +
|
|
"1715641054.367201\tCxT122\t10.0.0.2\t52.12.0.2\n" +
|
|
"1715641054.367201\tCxT122\t10.0.0.2\t52.12.0.2\n" +
|
|
"1715641114.367201\tCxT123\t10.0.0.3\t52.12.0.3\n" +
|
|
"1715641114.367201\tCxT123\t10.0.0.3\t52.12.0.3\n" +
|
|
"1715641174.367201\tCxT124\t10.0.0.4\t52.12.0.4\n" +
|
|
"1715641174.367201\tCxT124\t10.0.0.4\t52.12.0.4\n" +
|
|
"1715641234.367201\tCxT125\t10.0.0.5\t52.12.0.5\n" +
|
|
"1715641234.367201\tCxT125\t10.0.0.5\t52.12.0.5\n",
|
|
)
|
|
bytesWritten, err := file.Write(log)
|
|
require.NoError(t, err, "writing data to file should not produce an error")
|
|
require.Equal(t, len(log), bytesWritten, "number of bytes written should be equal to the length of the log data")
|
|
|
|
err = file.Close()
|
|
require.NoError(t, err, "closing file should not produce an error")
|
|
|
|
// import the mock data
|
|
_, err = cmd.RunImportCmd(time.Now(), d.cfg, afs, directory, "testDB", false, false)
|
|
require.NoError(t, err, "importing data should not produce an error")
|
|
|
|
// connect to the database
|
|
db, err := database.ConnectToDB(context.Background(), "testDB", d.cfg, nil)
|
|
require.NoError(t, err, "connecting to database should not produce an error")
|
|
|
|
var result struct {
|
|
Count uint64 `ch:"count"`
|
|
}
|
|
|
|
// create a query to count number of conn records
|
|
err = db.Conn.QueryRow(db.GetContext(), `
|
|
SELECT count() as count FROM conn
|
|
`).ScanStruct(&result)
|
|
require.NoError(t, err, "querying conn table should not produce an error")
|
|
require.Equal(t, uint64(10), result.Count, "number of records in conn table should be 10")
|
|
|
|
// get the min timestamp
|
|
minTS, _, _, err := db.GetBeaconMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching min timestamp should not produce an error")
|
|
|
|
// get the network size (number of unique internal hosts for the past 24 hours)
|
|
networkSize, err := db.GetNetworkSize(minTS)
|
|
require.NoError(t, err, "getting network size should not produce an error")
|
|
|
|
// verify the expected network size
|
|
fmt.Println("Network Size: ", networkSize)
|
|
require.Greater(t, networkSize, uint64(0), "network size should be greater than zero")
|
|
require.Equal(t, uint64(5), networkSize, "network size for test 2 should be match expected value")
|
|
|
|
// remove the mock directory
|
|
require.NoError(t, afs.RemoveAll(directory), "removing mock directory should not produce an error")
|
|
|
|
})
|
|
|
|
d.Run("Open Conn Dataset", func() {
|
|
t := d.T()
|
|
// import a dataset
|
|
_, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/open_conns/open/open_conn.log", "testDBzzzz", false, false)
|
|
require.NoError(t, err, "importing data should not produce an error")
|
|
|
|
// connect to the database
|
|
db, err := database.ConnectToDB(context.Background(), "testDBzzzz", d.cfg, nil)
|
|
require.NoError(t, err, "connecting to database should not produce an error")
|
|
|
|
var result struct {
|
|
Count uint64 `ch:"count"`
|
|
}
|
|
|
|
// validate that number of conn records is zero
|
|
err = db.Conn.QueryRow(db.GetContext(), `
|
|
SELECT count() as count FROM conn
|
|
`).ScanStruct(&result)
|
|
require.NoError(t, err, "querying conn table should not produce an error")
|
|
require.Equal(t, uint64(0), result.Count, "number of records in conn table should be 0")
|
|
|
|
// validate that number of open conn records is > 0
|
|
err = db.Conn.QueryRow(db.GetContext(), `
|
|
SELECT count() as count FROM openconn
|
|
`).ScanStruct(&result)
|
|
require.NoError(t, err, "querying openconn table should not produce an error")
|
|
require.Greater(t, result.Count, uint64(0), "number of open conn records should be >0")
|
|
|
|
// get the min timestamp
|
|
minTS, _, _, _, err := db.GetTrueMinMaxTimestamps()
|
|
require.NoError(t, err, "fetching min timestamp should not produce an error")
|
|
|
|
// get the network size (number of unique internal hosts for the past 24 hours)
|
|
networkSize, err := db.GetNetworkSize(minTS)
|
|
require.NoError(t, err, "getting network size should not produce an error")
|
|
|
|
// verify the expected network size
|
|
// fmt.Println("Network Size: ", networkSize)
|
|
require.Greater(t, networkSize, uint64(0), "network size should be greater than zero")
|
|
require.Equal(t, uint64(12), networkSize, "network size for test should match expected value")
|
|
|
|
})
|
|
|
|
}
|