activecm-rita/integration_rolling/ttl_test.go

package integration_rolling_test

import (
	"context"
	"fmt"
	"log"
	"os"
	"strconv"
	"testing"
	"time"

	"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/ClickHouse/clickhouse-go/v2"
	"github.com/joho/godotenv"
	"github.com/spf13/afero"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	tc "github.com/testcontainers/testcontainers-go/modules/compose"
)

// TearDownSuite is run once after all tests have finished
func (d *TTLTestSuite) TearDownSuite() {
	d.Require().NoError(os.RemoveAll(logDir))
	d.Require().NoError(os.RemoveAll(futureLogDir))
	d.cleanupContainer()
}

func SetupClickHouseTTL(t *testing.T) (time.Time, func(time.Duration) error, func()) {
	t.Helper()
	identifier := tc.StackIdentifier("ttl_testing")
	compose, err := tc.NewDockerComposeWith(tc.WithStackFiles("./test.docker-compose.yml"), identifier)
	require.NoError(t, err, "NewDockerComposeAPIWith()")

	cleanupContainer := func() {
		require.NoError(t, compose.Down(context.Background(), tc.RemoveOrphans(true), tc.RemoveVolumes(true), tc.RemoveImagesLocal), "compose.Down()")
	}

	ctx := context.Background()

	err = compose.Up(ctx, tc.Wait(true), tc.WithRecreate("RecreateForce"))
	require.NoError(t, err)

	require.NoError(t, err)

	importTime := time.Now().UTC()
	changeTimezone := func(dur time.Duration) error {
		container, err := compose.ServiceContainer(context.Background(), "clickhouse")
		require.NoError(t, err)
		s := importTime.Add(dur).Unix()
		status, _, err := container.Exec(context.Background(), []string{"date", "+%m-%d-%Y %H:%M:%S", "-s", fmt.Sprintf("@%d", s)})

		require.NoError(t, err)
		require.EqualValues(t, 0, status)
		return nil
	}
	return importTime, changeTimezone, cleanupContainer

}

type TTLTestSuite struct {
	suite.Suite
	cfg              *config.Config
	importTime       time.Time
	cleanupContainer func()
	changeTime       func(time.Duration) error
	server           *database.ServerConn
}

const ConfigPath = "../integration/test_config.hjson"

func TestTTLs(t *testing.T) {
	if err := godotenv.Overload("../.env", "./test.env"); err != nil {
		log.Fatal("Error loading .env file")
	}
	suite.Run(t, new(TTLTestSuite))
}

// SetupSubTest is run before the each subtest
func (d *TTLTestSuite) SetupSubTest() {
	t := d.T()
	// reset the time to the current time before each subtest
	require.NoError(t, d.changeTime(0))
}

// SetupSuite is run once before the first test starts
func (d *TTLTestSuite) SetupSuite() {
	t := d.T()

	// load the config file
	cfg, err := config.ReadFileConfig(afero.NewOsFs(), ConfigPath)
	require.NoError(t, err, "config should load without error")

	// start clickhouse container
	importTime, changeTimezone, cleanup := SetupClickHouseTTL(t)
	d.importTime = importTime
	// update the config to use the clickhouse container connection
	cfg.Env.DBConnection = "localhost:9001"
	d.changeTime = changeTimezone
	d.cleanupContainer = cleanup
	d.cfg = cfg

	// connect to clickhouse server
	server, err := database.ConnectToServer(context.Background(), d.cfg)
	require.NoError(t, err, "connecting to server should not produce an error")
	d.server = server

	// verify timezone (since we are going to be changing it to simulate the TTL process)
	var timezone string
	err = server.Conn.QueryRow(server.GetContext(), "SELECT serverTimeZone()").Scan(&timezone)
	require.NoError(t, err)
	require.Equal(t, "UTC", timezone, "timezone should be UTC before importing")

	err = server.Conn.Exec(server.GetContext(), "TRUNCATE TABLE IF EXISTS metadatabase.files")
	require.NoError(t, err)

	err = server.Conn.Exec(server.GetContext(), "TRUNCATE TABLE IF EXISTS metadatabase.min_max")
	require.NoError(t, err)

	err = server.Conn.Exec(server.GetContext(), "TRUNCATE TABLE IF EXISTS metadatabase.imports")
	require.NoError(t, err)

	err = server.Conn.Exec(server.GetContext(), "TRUNCATE TABLE IF EXISTS metadatabase.historical_first_seen")
	require.NoError(t, err)

}

// go test -count=1 -v ./integration_rolling -run TestTTLs/TestTableTTLs
func (d *TTLTestSuite) TestTableTTLs() {
	type importData struct {
		directory string

		importStartTime         time.Time
		age                     time.Duration
		expectedAgeOutAt26Hours bool
		expectedAgeOutAt2Weeks  bool
	}

	// all these tables will have a ttl interval of 26 hours
	// interval := 26 * time.Hour

	testCases := []struct {
		name    string
		afs     afero.Fs
		rolling bool
		imports []importData
	}{
		{
			name:    "Single Import - 26hrs",
			afs:     afero.NewOsFs(),
			rolling: true,
			imports: []importData{
				{

					directory:               "../test_data/valid_tsv",
					age:                     0,
					expectedAgeOutAt26Hours: true,
					expectedAgeOutAt2Weeks:  false,
				},
			},
		},
	}
	for index, test := range testCases {
		d.Run("Import: "+test.name, func() {
			t := d.T()

			dbName := "testDB" + strconv.Itoa(index) // Convert index to string

			// iterate over each rolling import
			for i := range test.imports {
				// set up test variables
				rebuild := i == 0
				importData := &test.imports[i] // pointer to modify the slice

				// set import start time
				importData.importStartTime = d.importTime.Add(importData.age)

				// fmt.Println("TEST TIME", importData.importStartTime, "ACTUAL", time.Now().UTC().Add(importData.age), importData.age, tc.buffer)

				// import the mock data
				_, err := cmd.RunImportCmd(importData.importStartTime, d.cfg, test.afs, importData.directory, dbName, test.rolling, rebuild)
				require.NoError(t, err, "importing data should not produce an error")
			}
		})
	}

	gT := d.T()
	// connect to the database
	metaDB, err := database.ConnectToDB(context.Background(), "metadatabase", d.cfg, nil)
	require.NoError(gT, err, "connecting to database should not produce an error")

	// verify that time is equal to now (unchanged)
	verifyTimeChange(gT, metaDB, 0, 10)

	// trigger table merge (TTL)
	fmt.Println("Triggering MetaDB table merges...")
	optimizeMetaDBTables(gT, metaDB, d.changeTime, 0, "")
	fmt.Println("Done merging.")

	// verify data was imported to log tables
	for index, test := range testCases {
		t := d.T()
		dbName := "testDB" + strconv.Itoa(index) // Convert index to string
		// connect to the database
		db, err := database.ConnectToDB(context.Background(), dbName, d.cfg, nil)
		require.NoError(gT, err, "connecting to database should not produce an error")
		fmt.Println("Triggering table merges...")
		optimizeTables(t, db)
		fmt.Println("Done merging.")

		for i := range test.imports {
			t.Run(fmt.Sprintf("post import data check: %s %d", dbName, i), func(t *testing.T) {
				if test.imports[i].age == 0 {
					log.Println("VERIFY IMPORT FOR", test.imports[i].importStartTime)
					verifyTables(t, db, test.imports[i].importStartTime, false, false, false, false)
				}
			})
		}
	}

	// change time to trigger the TTL process
	require.NoError(gT, d.changeTime(26*time.Hour), "changing time should not produce an error")

	// verify time change
	verifyTimeChange(d.T(), metaDB, 26*time.Hour, 10)

	fmt.Printf("\nChanged container time to 26 hours in the future\n")

	// trigger table merge (TTL)
	fmt.Println("Triggering MetaDB table merges...")
	optimizeMetaDBTables(gT, metaDB, d.changeTime, 26*time.Hour, "")
	fmt.Println("Done merging.")
	verifyTimeChange(d.T(), metaDB, 26*time.Hour, 10)

	// loop over imports
	for index, test := range testCases {
		t := d.T()
		dbName := "testDB" + strconv.Itoa(index) // Convert index to string
		// connect to the database
		db, err := database.ConnectToDB(context.Background(), dbName, d.cfg, nil)
		require.NoError(gT, err, "connecting to database should not produce an error")
		verifyTimeChange(t, db, 26*time.Hour, 10)

		fmt.Println("Triggering table merges...")
		optimizeTables(t, db)
		fmt.Println("Done merging.")
		verifyTimeChange(t, db, 26*time.Hour, 10)

		for i := range test.imports {
			t.Run(fmt.Sprintf("post +26h check %s %d", dbName, i), func(t *testing.T) {
				// check to see if data from old imports (>=2w old) are out of the dataset
				expected2wEmpty := test.imports[i].expectedAgeOutAt2Weeks
				verifyTables(t, db, test.imports[i].importStartTime, true, expected2wEmpty, false, false)
			})
		}
	}

	err = d.changeTime(14 * 24 * time.Hour) // 2 weeks
	require.NoError(gT, err, "changing time should not produce an error")

	// verify time change
	verifyTimeChange(gT, metaDB, ((time.Hour * 24) * 14), 10)
	fmt.Printf("\nChanged container time to 2 weeks in the future\n")

	// trigger table merge (TTL)
	fmt.Println("Triggering MetaDB table merges...")
	optimizeMetaDBTables(gT, metaDB, d.changeTime, 14*24*time.Hour, "")
	fmt.Println("Done merging.")

	for index, test := range testCases {
		t := d.T()
		dbName := "testDB" + strconv.Itoa(index) // Convert index to string
		// connect to the database
		db, err := database.ConnectToDB(context.Background(), dbName, d.cfg, nil)
		require.NoError(gT, err, "connecting to database should not produce an error")
		fmt.Println("Triggering table merges...")
		optimizeTables(t, db)
		fmt.Println("Done merging.")

		for i := range test.imports {
			t.Run(fmt.Sprintf("post +2 weeks check %d", i), func(t *testing.T) {
				verifyTables(t, db, test.imports[i].importStartTime, true, true, false, false)
			})
		}
	}

	err = d.changeTime(181 * 24 * time.Hour) // 6 months
	require.NoError(gT, err, "changing time should not produce an error")

	// verify time change
	verifyTimeChange(gT, metaDB, (181 * 24 * time.Hour), 10)
	fmt.Printf("\nChanged container time to 6 months in the future\n")

	// trigger table merge (TTL)
	fmt.Println("Triggering MetaDB table merges...")
	optimizeMetaDBTables(gT, metaDB, d.changeTime, 181*24*time.Hour, "files")
	fmt.Println("Done merging.")

	for index, test := range testCases {
		t := d.T()
		dbName := "testDB" + strconv.Itoa(index) // Convert index to string
		// connect to the database
		db, err := database.ConnectToDB(context.Background(), dbName, d.cfg, nil)
		require.NoError(gT, err, "connecting to database should not produce an error")
		for i := range test.imports {
			t.Run(fmt.Sprintf("post +6 months check %d", i), func(t *testing.T) {
				verifyTables(t, db, test.imports[i].importStartTime, true, true, true, false)
			})
		}
	}

	err = d.changeTime(366 * 24 * time.Hour) // 1 year
	require.NoError(gT, err, "changing time should not produce an error")

	// verify time change
	verifyTimeChange(gT, metaDB, (366 * 24 * time.Hour), 10)
	fmt.Printf("\nChanged container time to 1 year in the future\n")

	// trigger table merge (TTL)
	fmt.Println("Triggering MetaDB table merges...")
	optimizeMetaDBTables(gT, metaDB, d.changeTime, 366*24*time.Hour, "imports")
	fmt.Println("Done merging.")

	for index, test := range testCases {
		t := d.T()
		dbName := "testDB" + strconv.Itoa(index) // Convert index to string
		// connect to the database
		db, err := database.ConnectToDB(context.Background(), dbName, d.cfg, nil)
		require.NoError(gT, err, "connecting to database should not produce an error")
		for i := range test.imports {
			t.Run(fmt.Sprintf("post +1 year check %d", i), func(t *testing.T) {
				verifyMetaDBCountsByID(t, db, test.imports[i].importStartTime, []bool{true, true})
			})
		}
	}

}

// all tables should be optimized after each time the clock time has been changed
// regardless of whether or not they should be empty or not
func optimizeTables(t *testing.T, db *database.DB) {
	t.Helper()

	sensorTables := []string{"conn", "uconn", "http", "ssl", "usni", "dns", "udns", "pdns_raw", "pdns", "mime_type_uris",
		"threat_mixtape", "port_info", "http_proto", "tls_proto", "rare_signatures", "big_ol_histogram", "exploded_dns"}

	for _, table := range sensorTables {
		// require.NoError(t, d.changeTime("+26 hours"), "changing time should not produce an error")

		ctx := db.QueryParameters(clickhouse.Parameters{
			"database": db.GetSelectedDB(),
			"table":    table,
		})

		err := db.Conn.Exec(ctx, `OPTIMIZE TABLE {database:Identifier}.{table:Identifier} FINAL`)
		require.NoError(t, err, "optimizing %s.%s should not fail", db.GetSelectedDB(), table)
	}

}

func optimizeMetaDBTables(t *testing.T, db *database.DB, updateTime func(time.Duration) error, timeDelta time.Duration, justOne string) {
	t.Helper()
	metaDBTables := []string{"imports", "historical_first_seen", "files", "min_max", "valid_mime_types", "threat_intel", "threat_intel_feeds"}
	for _, table := range metaDBTables {
		// skip optimizing the other tables if justOne is set
		if justOne != "" {
			if table != justOne {
				continue
			}
		}
		ctx := db.QueryParameters(clickhouse.Parameters{
			"table": table,
		})
		err := updateTime(timeDelta)
		require.NoError(t, err)
		err = db.Conn.Exec(ctx, `OPTIMIZE TABLE metadatabase.{table:Identifier} FINAL`)
		require.NoError(t, err, "optimizing %s.%s should not fail", "metadatabase", table)
	}
}
func verifyTables(t *testing.T, db *database.DB, importTime time.Time, expect26hEmpty, expect2wEmpty, expect6mEmpty, expect1YrEmpty bool) {
	t.Helper()
	// verify the 24 hour log tables
	verifyLogTableCounts(t, db, importTime, expect26hEmpty)
	verifyLogTableHourViewsCounts(t, db, importTime, expect26hEmpty)
	verifyLogTableDayViewsCounts(t, db, importTime, expect26hEmpty)

	// verify the 2 week log tables
	verifyAnalysisSnapshotCounts(t, db, importTime, expect2wEmpty)
	verifyAnalysisSnapshotAnalyzedAtCounts(t, db, importTime, expect2wEmpty)

	// verify the 3 month log tables
	verifyMetaDBCountsByID(t, db, importTime, []bool{expect1YrEmpty, expect6mEmpty})
	// verifyMetaDBImportCounts(t, db, importTime, expect3mEmpty)
}

func verifyTimeChange(t *testing.T, db *database.DB, expectedDifference time.Duration, acceptableDelta int) {
	t.Helper()
	realNow := time.Now()
	var now time.Time
	err := db.Conn.QueryRow(db.GetContext(), "SELECT now()").Scan(&now)
	require.NoError(t, err)
	fmt.Println("RN", realNow.UTC().Unix(), "DB NOW", now.UTC().Unix(), "DIFF", realNow.UTC().Unix()-now.UTC().Unix())
	require.InDelta(t, realNow.UTC().Unix(), now.UTC().Unix(), float64(expectedDifference)+float64(acceptableDelta))
}

func verifyLogTableCounts(t *testing.T, db *database.DB, importTime time.Time, shouldBeEmpty bool) {
	t.Helper()
	query := `--sql
		SELECT count() FROM {table:Identifier}
		WHERE import_time=fromUnixTimestamp({importStartTime:Int64})
	`
	verifyTableCounts(t, db, strconv.FormatInt(importTime.Unix(), 10), shouldBeEmpty, database.LogTableTTLs, query)
}

func verifyLogTableHourViewsCounts(t *testing.T, db *database.DB, importTime time.Time, shouldBeEmpty bool) {
	t.Helper()
	query := `--sql
		SELECT count() FROM {table:Identifier}
		WHERE import_hour=toStartOfHour(fromUnixTimestamp({importStartTime:Int64}))
	`
	verifyTableCounts(t, db, strconv.FormatInt(importTime.Unix(), 10), shouldBeEmpty, database.LogTableViewsHourTTLs, query)
}

func verifyLogTableDayViewsCounts(t *testing.T, db *database.DB, importTime time.Time, shouldBeEmpty bool) {
	t.Helper()
	query := `--sql
		SELECT count() FROM {table:Identifier}
		WHERE import_day=toStartOfDay(fromUnixTimestamp({importStartTime:Int64}))
	`
	verifyTableCounts(t, db, strconv.FormatInt(importTime.Unix(), 10), shouldBeEmpty, database.LogTableViewsDayTTLs, query)
}

func verifyAnalysisSnapshotCounts(t *testing.T, db *database.DB, importTime time.Time, shouldBeEmpty bool) {
	t.Helper()
	query := `--sql
		SELECT count() FROM {table:Identifier}
		WHERE import_hour=toStartOfHour(fromUnixTimestamp({importStartTime:Int64}))
	`
	verifyTableCounts(t, db, strconv.FormatInt(importTime.Unix(), 10), shouldBeEmpty, database.AnalysisSnapshotHourTTLs, query)
}

func verifyAnalysisSnapshotAnalyzedAtCounts(t *testing.T, db *database.DB, importTime time.Time, shouldBeEmpty bool) {
	t.Helper()
	query := `--sql
		SELECT count() FROM {table:Identifier}
		WHERE analyzed_at = fromUnixTimestamp64Micro({importStartTime:Int64})
	`
	verifyTableCounts(t, db, strconv.FormatInt(importTime.UnixMicro(), 10), shouldBeEmpty, database.AnalysisSnapshotAnalyzedAtTTLs, query)
}

func verifyMetaDBCountsByID(t *testing.T, db *database.DB, importTime time.Time, shouldBeEmpty []bool) {
	t.Helper()
	importID, err := util.NewFixedStringHash(strconv.FormatInt(importTime.UnixMicro(), 10))
	require.NoError(t, err)

	tables := []string{"imports", "files"}
	for i, table := range tables {

		ctx := db.QueryParameters(clickhouse.Parameters{
			"importID": importID.Hex(),
			"table":    table,
		})
		query := `--sql
			SELECT count() FROM metadatabase.{table:Identifier}
			WHERE import_id = unhex({importID:String})
		`
		if table == "files" {
			query += `--sql
			AND rolling = true`
		}
		var count uint64
		err := db.Conn.QueryRow(ctx, query).Scan(&count)
		require.NoError(t, err, "querying table metadatabase.%sshould not produce an error", table)

		if shouldBeEmpty[i] {
			require.Equal(t, uint64(0), count, "table metadatabase.%s should be empty", table)
		} else {
			require.Greater(t, count, uint64(0), "table metadatabase.%s should have more than 0 rows", table)
		}
	}
}

func verifyTableCounts(t *testing.T, db *database.DB, importTime string, shouldBeEmpty bool, tables []string, query string) {
	t.Helper()
	for _, table := range tables {
		ctx := db.QueryParameters(clickhouse.Parameters{
			"table":           table,
			"importStartTime": importTime,
		})

		var count uint64
		err := db.Conn.QueryRow(ctx, query).Scan(&count)
		require.NoError(t, err, "querying table %s should not produce an error", table)

		if shouldBeEmpty {
			require.Equal(t, uint64(0), count, "table %s.%s should be empty", db.GetSelectedDB(), table)
		} else {
			require.Greater(t, count, uint64(0), "table %s.%s should have more than 0 rows for %s", db.GetSelectedDB(), table, importTime)
		}
	}
}