fixed linting errors, reworked config loading, removed unused code, fixed some tests

Co-Authored-By: Naomi Kramer <naomiagoddard@gmail.com>
This commit is contained in:
Liza Tsibur
2024-07-15 21:08:32 -06:00
parent c47ac8082b
commit b9974cea30
56 changed files with 924 additions and 973 deletions
+70 -10
View File
@@ -34,10 +34,10 @@ linters-settings:
sizeThreshold: 512 # Default: 512
skipTestFuncs: true
rangeValCopy:
sizeThreshold: 1000 # Default: 128
sizeThreshold: 200 # Default: 128
skipTestFuncs: true
hugeParam:
sizeThreshold: 80 # Default: 80
sizeThreshold: 100 # Default: 80
# Examines Go source code and reports suspicious constructs. It is roughly the same as 'go vet' and uses its passes.
govet:
@@ -76,8 +76,51 @@ linters-settings:
require-specific: true
revive:
enable-all-rules: true
# https://github.com/mgechev/revive/blob/master/RULES_DESCRIPTIONS.md
rules:
- name: add-constant
disabled: true
- name: argument-limit
arguments: [10]
- name: cognitive-complexity # enable in future cleanup
disabled: true
- name: confusing-results
disabled: true
- name: cyclomatic # enable in future cleanup
disabled: true
- name: deep-exit
disabled: true # enable in future cleanup
- name: empty-lines
disabled: true
- name: exported
arguments: ["disableStutteringCheck"]
- name: flag-parameter
disabled: true
- name: function-length
disabled: true
- name: function-result-limit
arguments: [8]
- name: get-return
disabled: true
- name: if-return
disabled: true
- name: line-length-limit
arguments: [250]
- name: max-control-nesting
arguments: [7]
- name: max-public-structs
arguments: [8]
- name: var-naming
arguments:
- [ ] # allowlist
- [ ] # blocklist
- [ { "upperCaseConst": true, "skipPackageNameChecks": true } ]
- name: unexported-return
disabled: true
- name: unhandled-error
arguments: [ "fmt.Printf", "fmt.Fprintf", "fmt.Println", "io.Closer.Close", "os.File.Close" ]
- name: unused-receiver
disabled: true
goconst:
@@ -162,7 +205,6 @@ linters:
# Both require a bit more explicit returns.
- nilerr
- nilnil
# Finds sending HTTP request without context.Context.
- noctx
@@ -176,9 +218,6 @@ linters:
# Lint your Prometheus metrics name
- promlinter
# Checks that package variables are not reassigned.
- reassign
# Drop-in replacement of `golint`.
- revive
@@ -221,8 +260,29 @@ linters:
- zerologlint
issues:
exclude:
- "var-naming: don't use an underscore in package name"
# exclude:
# - "var-naming: don't use an underscore in package name"
exclude-rules:
- text: 'shadow: declaration of "(err|ctx)" shadows declaration at'
linters: [ govet ]
# ignore shadowing of err and ctx
- linters: [ govet ]
text: 'shadow: declaration of "(err|ctx)" shadows declaration at'
# tolerate long functions in tests
- linters: [ revive ]
path: "(.+)_test.go"
text: "function-length: .*"
# tolerate long lines in tests
- linters: [ revive ]
path: "(.+)_test.go"
text: "line-length-limit: .*"
# tolerate deep-exit in tests
- linters: [ revive ]
path: "(.+)_test.go"
text: "deep-exit: .*"
# tolerate import shadowing in tests
- linters: [ revive ]
path: "(.+)_test.go"
text: "import-shadowing: The name ('require') shadows an import name"
+13 -20
View File
@@ -10,17 +10,13 @@ import (
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/database"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/util"
"golang.org/x/sync/errgroup"
"golang.org/x/time/rate"
)
func (r ThreatMixtape) FilterValue() string { return r.Src.String() }
func (i ThreatMixtape) Title() string { return i.Dst.String() }
func (i ThreatMixtape) Description() string { return fmt.Sprint(i.FinalScore * 100) }
type Analyzer struct {
Database *database.DB
ImportID util.FixedString
@@ -80,13 +76,6 @@ type ThreatMixtape struct {
FirstSeenScore float32 `ch:"first_seen_score"`
ThreatIntelDataSizeScore float32 `ch:"threat_intel_data_size_score"`
MissingHostHeaderScore float32 `ch:"missing_host_header_score"`
// FirstSeenMod float32 `ch:"first_seen_mod"`
// PrevalenceMod float32 `ch:"prevalence_mod"`
// MissingHostHeaderMod float32 `ch:"missing_host_header_mod"`
// C2OverDNSDirectConnMod float32 `ch:"c2_over_dns_direct_conn_mod"`
// MIMETypeMismatchMod float32 `ch:"mime_type_mismatch_mod"`
// RareSignatureMod float32 `ch:"rare_signature_mod"`
}
// NewAnalyzer returns a new Analyzer object
@@ -100,10 +89,6 @@ func NewAnalyzer(db *database.DB, cfg *config.Config, importID util.FixedString,
}
var firstSeenMaxTS time.Time
if !useCurrentTime {
// _, firstSeenMaxTS, _, _, err = db.GetTrueMinMaxTimestamps()
// if err != nil {
// return nil, fmt.Errorf("could not get max timestamp for first seen analysis: %w", err)
// }
firstSeenMaxTS = maxTS
}
@@ -128,7 +113,7 @@ func NewAnalyzer(db *database.DB, cfg *config.Config, importID util.FixedString,
}
func (analyzer *Analyzer) Analyze() error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// log the start time of the analysis
start := time.Now()
@@ -174,7 +159,7 @@ func (analyzer *Analyzer) Analyze() error {
}
func (analyzer *Analyzer) runAnalysis() error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// loop over the uconn channel to process each entry
for entry := range analyzer.UconnChan {
@@ -189,7 +174,11 @@ func (analyzer *Analyzer) runAnalysis() error {
// set the first seen historical value
firstSeenHistorical, replaced := util.ValidateTimestamp(entry.FirstSeenHistorical)
if replaced {
logger.Debug().Str("src", entry.Src.String()).Str("dst", entry.Dst.String()).Str("missing_host_count", fmt.Sprint(entry.MissingHostCount)).Str("fqdn", entry.FQDN).Msg("historical first seen timestamp was missing")
logger.Debug().
Str("src", entry.Src.String()).
Str("dst", entry.Dst.String()).
Str("missing_host_count", fmt.Sprint(entry.MissingHostCount)).
Str("fqdn", entry.FQDN).Msg("historical first seen timestamp was missing")
}
// if the last seen timestamp was not valid, then this entry cannot be inserted into the mixtape
@@ -197,7 +186,11 @@ func (analyzer *Analyzer) runAnalysis() error {
// this should log a warning as this is a bugs
lastSeen, replaced := util.ValidateTimestamp(entry.LastSeen)
if replaced {
logger.Debug().Str("src", entry.Src.String()).Str("dst", entry.Dst.String()).Str("missing_host_count", fmt.Sprint(entry.MissingHostCount)).Str("fqdn", entry.FQDN).Msg("last seen timestamp was missing")
logger.Debug().
Str("src", entry.Src.String()).
Str("dst", entry.Dst.String()).
Str("missing_host_count", fmt.Sprint(entry.MissingHostCount)).
Str("fqdn", entry.FQDN).Msg("last seen timestamp was missing")
}
mixtape.FirstSeenHistorical = firstSeenHistorical
+19 -5
View File
@@ -7,7 +7,7 @@ import (
"slices"
"sort"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/util"
"github.com/montanaflynn/stats"
@@ -31,7 +31,7 @@ type Beacon struct {
}
func (analyzer *Analyzer) analyzeBeacon(entry *AnalysisResult) (Beacon, error) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
var beacon Beacon
// verify that minTSBeacon < maxTSBeacon
@@ -55,14 +55,20 @@ func (analyzer *Analyzer) analyzeBeacon(entry *AnalysisResult) (Beacon, error) {
}
// calculate histogram score (note: we currently look at a 24 hour period)
_, _, totalBars, longestRun, histScore, err := getHistogramScore(analyzer.minTSBeacon.Unix(), analyzer.maxTSBeacon.Unix(), entry.TSList, analyzer.Config.Scoring.Beacon.HistModeSensitivity, analyzer.Config.Scoring.Beacon.HistBimodalOutlierRemoval, analyzer.Config.Scoring.Beacon.HistBimodalMinHours, 24)
_, _, totalBars, longestRun, histScore, err := getHistogramScore(
analyzer.minTSBeacon.Unix(), analyzer.maxTSBeacon.Unix(), entry.TSList, analyzer.Config.Scoring.Beacon.HistModeSensitivity,
analyzer.Config.Scoring.Beacon.HistBimodalOutlierRemoval, analyzer.Config.Scoring.Beacon.HistBimodalMinHours, 24,
)
if err != nil {
logger.Err(err).Caller().Str("src", entry.Src.String()).Str("dst", entry.Dst.String()).Str("fqdn", entry.FQDN).Send()
return beacon, err
}
// calculate duration score
_, _, durScore, err := getDurationScore(analyzer.minTSBeacon.Unix(), analyzer.maxTSBeacon.Unix(), int64(entry.TSList[0]), int64(entry.TSList[len(entry.TSList)-1]), totalBars, longestRun, analyzer.Config.Scoring.Beacon.DurMinHours, analyzer.Config.Scoring.Beacon.DurIdealNumberOfConsistentHours)
_, _, durScore, err := getDurationScore(
analyzer.minTSBeacon.Unix(), analyzer.maxTSBeacon.Unix(), int64(entry.TSList[0]), int64(entry.TSList[len(entry.TSList)-1]),
totalBars, longestRun, analyzer.Config.Scoring.Beacon.DurMinHours, analyzer.Config.Scoring.Beacon.DurIdealNumberOfConsistentHours,
)
if err != nil {
logger.Err(err).Caller().Str("src", entry.Src.String()).Str("dst", entry.Dst.String()).Str("fqdn", entry.FQDN).Send()
return beacon, err
@@ -129,6 +135,10 @@ func getBeaconScore(tsScore, tsWeight, dsScore, dsWeight, durScore, durWeight, h
return score, nil
}
// getTimestampScore calculates the timestamp score for a given list of timestamps. This score is based on the
// statistical properties of the intervals between timestamps, utilizing skewness and median absolute deviation
// to calculate a score that reflects the consistency of the intervals. This function returns the ts score, skew,
// median absolute deviation, intervals between timestamps, their counts, the most frequent interval, and its count.
func getTimestampScore(tsList []uint32) (float64, float64, float64, []int64, []int64, int64, int64, error) {
// ensure that the input slice has at least 4 elements (need at least 3 intervals, which requires at least 4 timestamps)
if len(tsList) < 4 {
@@ -184,6 +194,10 @@ func getTimestampScore(tsList []uint32) (float64, float64, float64, []int64, []i
}
// getDataSizeScore calculates the data size score for a given list of data sizes. This score is based on the
// statistical properties of the data sizes, utilizing skewness and median absolute deviation to calculate a
// score that reflects the consistency of the data sizes. This function returns the ds score, skew,
// median absolute deviation, unique data sizes, their counts, the most frequent data size, and its count.
func getDataSizeScore(bytesList []float64) (float64, float64, float64, []int64, []int64, int64, int64, error) {
// ensure that the input slice has at least 3 elements
if len(bytesList) < 3 {
@@ -209,7 +223,7 @@ func getDataSizeScore(bytesList []float64) (float64, float64, float64, []int64,
}
// calculateStatisticalScore calculates the statistical score, skew, and median absolute derivation for a given list of float64 values
// calculateStatisticalScore calculates the statistical score, skew, and median absolute deviation for a given list of float64 values
func calculateStatisticalScore(values []float64, defaultMadScore float64) (float64, float64, float64, error) {
// ensure that the input slice is not empty
if len(values) == 0 {
+25 -8
View File
@@ -293,6 +293,31 @@ func TestGetTimestampScore(t *testing.T) {
expectedScore: 0.995, // score = round(((1 + 0.99)/2)*1000)/1000 = 0.995
expectedError: false,
},
{
name: "Connection with Bi-Modal Intervals",
// timestamps : 1517338924, 1517338924 + 98, 1517338924 + 98 + 300, 1517338924 + 98 + 300 + 98, 1517338924 + 98 + 300 + 98 + 300, 1517338924 + 98 + 300 + 98 + 300 + 98, 1517338924 + 98 + 300 + 98 + 300 + 98 + 300, 1517338924 + 98 + 300 + 98 + 300 + 98 + 300 + 98, 1517338924 + 98 + 300 + 98 + 300 + 98 + 300 + 98 + 300, 1517338924 + 98 + 300 + 98 + 300 + 98 + 300 + 98 + 300 + 98,
tsList: []uint32{1517338924, 1517339022, 1517339322, 1517339420, 1517339720, 1517339818, 1517340118, 1517340216, 1517340516, 1517340614, 1517340914},
// intervals between timestamps: 98, 300, 98, 300, 98, 300, 98, 300, 98
expectedUniqueIntervals: []int64{98, 300},
expectedUniqueIntervalCounts: []int64{5, 5},
expectedTSMode: 98,
expectedTSModeCount: 5,
// Quartiles: {98 199 300}
// q1 = 98, q2 = 199, q3 = 300
// numerator = q3 + q1 - 2*q2 = 300 + 98 - 2*199 = 300 + 98 - 398 = 0
// IQR = q3 - q1 = 300 - 98 = 202
// Bowley Skewness = (Q3+Q1 2Q2) / (Q3 Q1) = numerator / IQR
// but if the denominator less than 10 or the median is equal to the lower or upper quartile, the skewness is zero
// skewness = 0
// skewness score = 1 - skewness = 1 - 0 = 1
// median absolute deviation = median(abs(x - median(x))) = median(abs(x - 199)) = median(101, 100, 101, 100, 101, 100, 101, 100, 101) = 101
// MAD score = 1 - (mad/median) = 1 - 101/199 = 0.492
// score = round(((1 + 0.492)/2)*1000)/1000 = 0.746
expectedSkew: 0,
expectedMAD: 101,
expectedScore: 0.746,
expectedError: false,
},
{
name: "Connection with Random Intervals",
tsList: []uint32{1517338924, 1517338925, 1517339224, 1517339249, 1517344224, 1517344314, 1517344316, 1517344358, 1517344858, 1517346358},
@@ -463,14 +488,6 @@ func TestGetDataSizeScore(t *testing.T) {
// run the function
score, skew, mad, sizes, sizeCounts, mode, modeCount, err := getDataSizeScore(test.bytesList)
// fmt.Println("score: ", score)
// fmt.Println("skew: ", skew)
// fmt.Println("mad: ", mad)
// fmt.Println("sizes: ", sizes)
// fmt.Println("sizeCounts: ", sizeCounts)
// fmt.Println("mode: ", mode)
// fmt.Println("modeCount: ", modeCount)
// fmt.Println("err: ", err)
// check if an error was expected
require.Equal(test.expectedError, err != nil, "Expected error to be %v, got %v", test.expectedError, err)
+17 -19
View File
@@ -4,11 +4,10 @@ import (
"context"
"fmt"
"net"
"os"
"strconv"
"time"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/progressbar"
"github.com/activecm/rita/v5/util"
@@ -61,7 +60,7 @@ type AnalysisResult struct {
}
func (analyzer *Analyzer) Spagoop(ctx context.Context) error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// record start time
start := time.Now()
@@ -117,8 +116,7 @@ func (analyzer *Analyzer) Spagoop(ctx context.Context) error {
queryGroup.Go(func() error {
_, err := bars.Run()
if err != nil {
fmt.Println("error running program:", err)
os.Exit(1)
logger.Error().Err(err).Msg("error running program")
}
return err
})
@@ -139,8 +137,8 @@ func (analyzer *Analyzer) Spagoop(ctx context.Context) error {
return nil
}
func (analyzer *Analyzer) ScoopSNIConns(ctx context.Context, progress *tea.Program) error {
logger := logger.GetLogger()
func (analyzer *Analyzer) ScoopSNIConns(ctx context.Context, bars *tea.Program) error {
logger := zlog.GetLogger()
// initialize progress bar variables
var totalSNI uint64
@@ -353,18 +351,18 @@ func (analyzer *Analyzer) ScoopSNIConns(ctx context.Context, progress *tea.Progr
// send the unique sni connections to the uconn analysis channel
analyzer.UconnChan <- res
if i%1000 == 0 {
progress.Send(progressbar.ProgressMsg{ID: 1, Percent: float64(i / totalSNI)})
bars.Send(progressbar.ProgressMsg{ID: 1, Percent: float64(i / totalSNI)})
}
i++
}
}
rows.Close()
progress.Send(progressbar.ProgressMsg{ID: 1, Percent: 1})
bars.Send(progressbar.ProgressMsg{ID: 1, Percent: 1})
return nil
}
func (analyzer *Analyzer) ScoopIPConns(ctx context.Context, progress *tea.Program) error {
logger := logger.GetLogger()
func (analyzer *Analyzer) ScoopIPConns(ctx context.Context, bars *tea.Program) error {
logger := zlog.GetLogger()
totalRows := uint64(0)
hasSetTotal := false
@@ -373,15 +371,15 @@ func (analyzer *Analyzer) ScoopIPConns(ctx context.Context, progress *tea.Progra
if !hasSetTotal {
totalRows = p.Rows
if totalRows == 0 {
progress.Send(progressbar.ProgressMsg{ID: 2, Percent: 1})
bars.Send(progressbar.ProgressMsg{ID: 2, Percent: 1})
}
hasSetTotal = true
} else {
// update the progress bar
if totalRows > 0 {
progress.Send(progressbar.ProgressMsg{ID: 2, Percent: float64((totalRows - p.Rows) / totalRows)})
bars.Send(progressbar.ProgressMsg{ID: 2, Percent: float64((totalRows - p.Rows) / totalRows)})
}
progress.Send(progressbar.ProgressMsg{ID: 2, Percent: 1})
bars.Send(progressbar.ProgressMsg{ID: 2, Percent: 1})
}
}), clickhouse.WithParameters(clickhouse.Parameters{
// use minTSBeacon because all entries in conn are used in beaconing and openconn data is not limited by the hour since the tables are truncated before each import
@@ -574,8 +572,8 @@ func (analyzer *Analyzer) ScoopIPConns(ctx context.Context, progress *tea.Progra
return nil
}
func (analyzer *Analyzer) ScoopDNS(ctx context.Context, progress *tea.Program) error {
logger := logger.GetLogger()
func (analyzer *Analyzer) ScoopDNS(ctx context.Context, bars *tea.Program) error {
logger := zlog.GetLogger()
totalRows := uint64(0)
hasSetTotal := false
@@ -586,15 +584,15 @@ func (analyzer *Analyzer) ScoopDNS(ctx context.Context, progress *tea.Program) e
if !hasSetTotal {
totalRows = p.Rows
if totalRows == 0 {
progress.Send(progressbar.ProgressMsg{ID: 3, Percent: 1})
bars.Send(progressbar.ProgressMsg{ID: 3, Percent: 1})
}
hasSetTotal = true
} else {
// update the progress bar
if totalRows > 0 {
progress.Send(progressbar.ProgressMsg{ID: 3, Percent: float64((totalRows - p.Rows) / totalRows)})
bars.Send(progressbar.ProgressMsg{ID: 3, Percent: float64((totalRows - p.Rows) / totalRows)})
}
progress.Send(progressbar.ProgressMsg{ID: 3, Percent: 1})
bars.Send(progressbar.ProgressMsg{ID: 3, Percent: 1})
}
}), clickhouse.WithParameters(clickhouse.Parameters{
+1 -1
View File
@@ -50,7 +50,7 @@ func CheckForUpdate(cCtx *cli.Context, afs afero.Fs) error {
}
// check for update if version is set
if cfg.UpdateCheckEnabled && currentVersion != "" {
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)
+3
View File
@@ -60,6 +60,9 @@ func (c *CmdTestSuite) SetupSuite() {
cfg, err := config.ReadFileConfig(afs, ConfigPath)
require.NoError(t, err, "config should load without error")
// // set version
// config.Version = ""
// start clickhouse container
c.SetupClickHouse(t)
+11 -9
View File
@@ -64,8 +64,15 @@ var DeleteCommand = &cli.Command{
if cCtx.Bool("non-interactive") {
prompt = false
}
// load config file
cfg, err := config.ReadFileConfig(afs, cCtx.String("config"))
if err != nil {
return err
}
// run the delete command
if err := runDeleteCmd(afs, cCtx.String("config"), input, trimmedName, prompt); err != nil {
if err := RunDeleteCmd(cfg, input, trimmedName, prompt); err != nil {
return err
}
@@ -78,24 +85,20 @@ var DeleteCommand = &cli.Command{
},
}
func runDeleteCmd(afs afero.Fs, configPath string, entry string, trimmedName string, ask bool) error {
func RunDeleteCmd(cfg *config.Config, entry string, trimmedName string, ask bool) error {
// validate the trimmed name
if len(trimmedName) == 0 {
return ErrTrimmedNameEmpty
}
// load config file
cfg, err := config.ReadFileConfig(afs, configPath)
if err != nil {
return err
}
// connect to server
server, err := database.ConnectToServer(context.Background(), cfg)
if err != nil {
return err
}
// set uo prompt for confirmation
prompt := promptui.Prompt{
Label: "Delete Dataset",
IsConfirm: true,
@@ -147,7 +150,6 @@ func runDeleteCmd(afs afero.Fs, configPath string, entry string, trimmedName str
}
fmt.Println("Successfully deleted dataset if it existed.")
}
return nil
+149 -47
View File
@@ -24,38 +24,41 @@ func (c *CmdTestSuite) TestDeleteCommand() {
expectedRemainingDbs []string
expectedError error
}{
{
name: "No Wildcards - Database Matching Trimmed Name Exactly",
args: []string{"app", "delete", "--ni", "bingbong"},
dbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
expectedDeletedDbs: []string{"bingbong"},
expectedRemainingDbs: []string{"prefix_bingbong", "bingbong123", "prefix_bingbong123"},
expectedError: nil,
},
{
name: "Prefix Wildcard - Databases Ending with Trimmed Name",
args: []string{"app", "delete", "--ni", "*bingbong"},
dbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
expectedDeletedDbs: []string{"bingbong", "prefix_bingbong"},
expectedRemainingDbs: []string{"bingbong123", "prefix_bingbong123"},
expectedError: nil,
},
{
name: "Suffix Wildcard - Databases Starting with Trimmed Name",
args: []string{"app", "delete", "--ni", "bingbong*"},
dbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
expectedDeletedDbs: []string{"bingbong", "bingbong123"},
expectedRemainingDbs: []string{"prefix_bingbong", "prefix_bingbong123"},
expectedError: nil,
},
{
name: "Both Wildcards - Databases Containing Trimmed Name",
args: []string{"app", "delete", "--ni", "*bingbong*"},
dbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
expectedDeletedDbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
expectedRemainingDbs: []string{},
expectedError: nil,
},
// TODO: in order to check the actual deletion step from calling the command, we need a way to
// pass in or use the test clickhouse DBConnection, right now it gets overridden when load config gets called
// {
// name: "No Wildcards - Database Matching Trimmed Name Exactly",
// args: []string{"app", "delete", "--ni", "--config=../config.hjson", "bingbong"},
// dbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
// expectedDeletedDbs: []string{"bingbong"},
// expectedRemainingDbs: []string{"prefix_bingbong", "bingbong123", "prefix_bingbong123"},
// expectedError: nil,
// },
// {
// name: "Prefix Wildcard - Databases Ending with Trimmed Name",
// args: []string{"app", "delete", "--ni", "--config=../config.hjson", "*bingbong"},
// dbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
// expectedDeletedDbs: []string{"bingbong", "prefix_bingbong"},
// expectedRemainingDbs: []string{"bingbong123", "prefix_bingbong123"},
// expectedError: nil,
// },
// {
// name: "Suffix Wildcard - Databases Starting with Trimmed Name",
// args: []string{"app", "delete", "--ni", "--config=../config.hjson", "bingbong*"},
// dbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
// expectedDeletedDbs: []string{"bingbong", "bingbong123"},
// expectedRemainingDbs: []string{"prefix_bingbong", "prefix_bingbong123"},
// expectedError: nil,
// },
// {
// name: "Both Wildcards - Databases Containing Trimmed Name",
// args: []string{"app", "delete", "--ni", "--config=../config.hjson", "*bingbong*"},
// dbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
// expectedDeletedDbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
// expectedRemainingDbs: []string{},
// expectedError: nil,
// },
{
name: "Too Many Arguments",
args: []string{"app", "delete", "dbname", "extra"},
@@ -75,12 +78,6 @@ func (c *CmdTestSuite) TestDeleteCommand() {
// create a new app and context
app, ctx := setupTestApp(commands, flags)
// import new databases
for _, db := range test.dbs {
_, err := cmd.RunImportCmd(time.Now(), c.cfg, afero.NewOsFs(), "../test_data/open_conns/open", db, false, false)
require.NoError(err, "importing data should not produce an error")
}
// run app with test.args
err := app.RunContext(ctx, test.args)
if test.expectedError != nil {
@@ -90,24 +87,127 @@ func (c *CmdTestSuite) TestDeleteCommand() {
require.NoError(err, "error should be nil")
}
// validate that the expected databases were deleted
})
}
}
func (c *CmdTestSuite) TestRunDeleteCmd() {
type importDB struct {
name string
logDir string
}
tests := []struct {
name string
entry string
afs afero.Fs
dbs []importDB
expectedDeletedDbs []string
expectedRemainingDbs []string
expectedError error
}{
{
name: "No Wildcards - Database Matching Trimmed Name Exactly",
entry: "bingbong",
afs: afero.NewOsFs(),
dbs: []importDB{
{"bingbong", "../test_data/open_conns/open"},
{"prefix_bingbong", "../test_data/open_conns/open"},
{"bingbong123", "../test_data/open_conns/open"},
{"prefix_bingbong123", "../test_data/open_conns/open"},
},
expectedDeletedDbs: []string{"bingbong"},
expectedRemainingDbs: []string{"prefix_bingbong", "bingbong123", "prefix_bingbong123"},
},
{
name: "Prefix Wildcard - Databases Ending with Trimmed Name",
entry: "*bingbong",
afs: afero.NewOsFs(),
dbs: []importDB{
{"bingbong", "../test_data/open_conns/open"},
{"prefix_bingbong", "../test_data/open_conns/open"},
{"bingbong123", "../test_data/open_conns/open"},
{"prefix_bingbong123", "../test_data/open_conns/open"},
},
expectedDeletedDbs: []string{"bingbong", "prefix_bingbong"},
expectedRemainingDbs: []string{"bingbong123", "prefix_bingbong123"},
expectedError: nil,
},
{
name: "Suffix Wildcard - Databases Starting with Trimmed Name",
entry: "bingbong*",
afs: afero.NewOsFs(),
dbs: []importDB{
{"bingbong", "../test_data/open_conns/open"},
{"prefix_bingbong", "../test_data/open_conns/open"},
{"bingbong123", "../test_data/open_conns/open"},
{"prefix_bingbong123", "../test_data/open_conns/open"},
},
expectedDeletedDbs: []string{"bingbong", "bingbong123"},
expectedRemainingDbs: []string{"prefix_bingbong", "prefix_bingbong123"},
expectedError: nil,
},
{
name: "Both Wildcards - Databases Containing Trimmed Name",
entry: "*bingbong*",
afs: afero.NewOsFs(),
dbs: []importDB{
{"bingbong", "../test_data/open_conns/open"},
{"prefix_bingbong", "../test_data/open_conns/open"},
{"bingbong123", "../test_data/open_conns/open"},
{"prefix_bingbong123", "../test_data/open_conns/open"},
},
expectedDeletedDbs: []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"},
expectedRemainingDbs: []string{},
expectedError: nil,
},
}
for _, test := range tests {
c.Run(test.name, func() {
t := c.T()
importStartedAt := time.Now()
// import all dbs
for _, db := range test.dbs {
importResults, err := cmd.RunImportCmd(importStartedAt, c.cfg, test.afs, db.logDir, db.name, false, true)
require.NoError(t, err, "running import command should not produce an error")
require.NotNil(t, importResults, "import results should not be nil")
}
// trim leading and trailing wildcards
trimmedName, err := cmd.TrimWildcards(test.entry)
require.NoError(t, err, "trimming wildcards should not produce an error")
// validate the trimmed name
err = cmd.ValidateDatabaseName(trimmedName)
require.NoError(t, err, "validating database name should not produce an error")
// run the delete command
err = cmd.RunDeleteCmd(c.cfg, test.entry, trimmedName, false)
if test.expectedError != nil {
require.Contains(t, err.Error(), test.expectedError.Error(), "error should contain expected value")
} else {
require.NoError(t, err, "error should be nil")
}
// get list of import databases
dbs, err := c.server.ListImportDatabases()
require.NoError(t, err, "listing databases should not produce an error")
dbString := database.GetFlatDatabaseList(dbs)
require.NoError(err, "listing databases should not produce an error")
// validate that the expected databases were deleted
for _, db := range test.expectedDeletedDbs {
require.NotContains(dbString, db, "database %s should have been deleted", db)
require.NotContains(t, dbString, db, "database %s should have been deleted", db)
}
// validate that the expected databases remain
for _, db := range test.expectedRemainingDbs {
require.Contains(dbString, db, "database %s should not have been deleted", db)
require.Contains(t, dbString, db, "database %s should not have been deleted", db)
}
})
}
}
func TestTrimWildcards(t *testing.T) {
@@ -135,14 +235,16 @@ func TestTrimWildcards(t *testing.T) {
}
}
//nolint:unused // will be used in the future
func validateCommandsExist(t *testing.T, commands []*cli.Command, expected []string) {
t.Helper()
expectedCmds := make(map[string]bool)
for _, expectedCmd := range expected {
expectedCmds[expectedCmd] = false
}
for _, cmd := range commands {
if _, ok := expectedCmds[cmd.Name]; ok {
expectedCmds[cmd.Name] = true
for _, command := range commands {
if _, ok := expectedCmds[command.Name]; ok {
expectedCmds[command.Name] = true
}
}
for expectedSubCmd, present := range expectedCmds {
+32 -34
View File
@@ -17,9 +17,9 @@ import (
"github.com/activecm/rita/v5/analysis"
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/database"
"github.com/activecm/rita/v5/importer"
"github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/modifier"
i "github.com/activecm/rita/v5/importer"
zlog "github.com/activecm/rita/v5/logger"
m "github.com/activecm/rita/v5/modifier"
"github.com/activecm/rita/v5/util"
"github.com/spf13/afero"
@@ -125,7 +125,7 @@ type ImportTimestamps struct {
maxTSBeacon time.Time
}
type ImportResults struct {
importer.ResultCounts
i.ResultCounts
ImportID []util.FixedString
ImportTimestamps []ImportTimestamps
}
@@ -133,7 +133,7 @@ type ImportResults struct {
func RunImportCmd(startTime time.Time, cfg *config.Config, afs afero.Fs, logDir string, dbName string, rolling bool, rebuild bool) (ImportResults, error) {
var importResults ImportResults
logger := logger.GetLogger()
logger := zlog.GetLogger()
// keep track of the cumulative elapsed time
importStartedAt := startTime
@@ -195,24 +195,22 @@ func RunImportCmd(startTime time.Time, cfg *config.Config, afs afero.Fs, logDir
continue
}
// log the start of the import
logger.Debug().Str("started_at", importStartedAt.String()).Msg(fmt.Sprintf("Importing hour %d/%d", hour+1, len(hourlyLogs)))
// reset temporary tables
err = db.ResetTemporaryTables()
if err != nil {
return importResults, err
}
// parse logs
// importStart := importStartedAt
// if hour > 0 || day > 0 {
// // add the duration of all imports up to now to the original importStartedAt date
// importStart = importStartedAt.Add(time.Duration(elapsedTime) * time.Nanosecond)
// }
importer, err := importer.NewImporter(db, cfg, importStartedAt, numDigesters, numParsers, numWriters)
// set up new importer
importer, err := i.NewImporter(db, cfg, importStartedAt, numDigesters, numParsers, numWriters)
if err != nil {
return importResults, err
}
// import the data
err = importer.Import(afs, files)
if err != nil {
return importResults, err
@@ -265,7 +263,7 @@ func RunImportCmd(startTime time.Time, cfg *config.Config, afs afero.Fs, logDir
}
// set up new modifier
modifier, err := modifier.NewModifier(db, cfg, importer.ImportID, minTS, maxTS)
modifier, err := m.NewModifier(db, cfg, importer.ImportID, minTS)
if err != nil {
return importResults, err
}
@@ -362,7 +360,7 @@ func parseFolderDate(folder string) (time.Time, error) {
// path of each regular file on the string channel. It sends the result of the
// walk on the error channel. If done is closed, WalkFiles abandons its work.
func WalkFiles(afs afero.Fs, root string) ([]HourlyZeekLogs, []WalkError, error) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// check if root is a valid directory or file
err := util.ValidateDirectory(afs, root)
@@ -464,20 +462,20 @@ func WalkFiles(afs afero.Fs, root string) ([]HourlyZeekLogs, []WalkError, error)
// check if the file is one of the accepted log types
var prefix string
switch {
case strings.HasPrefix(filepath.Base(path), importer.ConnPrefix) && !strings.HasPrefix(filepath.Base(path), importer.ConnSummaryPrefixUnderscore) && !strings.HasPrefix(filepath.Base(path), importer.ConnSummaryPrefixHyphen):
prefix = importer.ConnPrefix
case strings.HasPrefix(filepath.Base(path), importer.OpenConnPrefix):
prefix = importer.OpenConnPrefix
case strings.HasPrefix(filepath.Base(path), importer.DNSPrefix):
prefix = importer.DNSPrefix
case strings.HasPrefix(filepath.Base(path), importer.HTTPPrefix):
prefix = importer.HTTPPrefix
case strings.HasPrefix(filepath.Base(path), importer.OpenHTTPPrefix):
prefix = importer.OpenHTTPPrefix
case strings.HasPrefix(filepath.Base(path), importer.SSLPrefix):
prefix = importer.SSLPrefix
case strings.HasPrefix(filepath.Base(path), importer.OpenSSLPrefix):
prefix = importer.OpenSSLPrefix
case strings.HasPrefix(filepath.Base(path), i.ConnPrefix) && !strings.HasPrefix(filepath.Base(path), i.ConnSummaryPrefixUnderscore) && !strings.HasPrefix(filepath.Base(path), i.ConnSummaryPrefixHyphen):
prefix = i.ConnPrefix
case strings.HasPrefix(filepath.Base(path), i.OpenConnPrefix):
prefix = i.OpenConnPrefix
case strings.HasPrefix(filepath.Base(path), i.DNSPrefix):
prefix = i.DNSPrefix
case strings.HasPrefix(filepath.Base(path), i.HTTPPrefix):
prefix = i.HTTPPrefix
case strings.HasPrefix(filepath.Base(path), i.OpenHTTPPrefix):
prefix = i.OpenHTTPPrefix
case strings.HasPrefix(filepath.Base(path), i.SSLPrefix):
prefix = i.SSLPrefix
case strings.HasPrefix(filepath.Base(path), i.OpenSSLPrefix):
prefix = i.OpenSSLPrefix
default: // skip file if it doesn't match any of the accepted prefixes
walkErrors = append(walkErrors, WalkError{Path: path, Error: ErrInvalidLogType})
continue
@@ -520,17 +518,17 @@ func WalkFiles(afs afero.Fs, root string) ([]HourlyZeekLogs, []WalkError, error)
for hour := range logMap[day] {
// if there are no conn logs in the hour, we have to skip any SSL and HTTP logs for that hour
if len(logMap[day][hour][importer.ConnPrefix]) == 0 && (len(logMap[day][hour][importer.SSLPrefix]) > 0 || len(logMap[day][hour][importer.HTTPPrefix]) > 0) {
if len(logMap[day][hour][i.ConnPrefix]) == 0 && (len(logMap[day][hour][i.SSLPrefix]) > 0 || len(logMap[day][hour][i.HTTPPrefix]) > 0) {
logger.Warn().Msg("SSL / HTTP logs are present, but no conn logs exist, skipping SSL / HTTP logs...")
delete(logMap[day][hour], importer.SSLPrefix)
delete(logMap[day][hour], importer.HTTPPrefix)
delete(logMap[day][hour], i.SSLPrefix)
delete(logMap[day][hour], i.HTTPPrefix)
}
// // if there are no open conn logs in the hour, we have to skip any open SSL and open HTTP logs for that hour
if len(logMap[day][hour][importer.OpenConnPrefix]) == 0 && (len(logMap[day][hour][importer.OpenSSLPrefix]) > 0 || len(logMap[day][hour][importer.OpenHTTPPrefix]) > 0) {
if len(logMap[day][hour][i.OpenConnPrefix]) == 0 && (len(logMap[day][hour][i.OpenSSLPrefix]) > 0 || len(logMap[day][hour][i.OpenHTTPPrefix]) > 0) {
logger.Warn().Msg("Open SSL / open HTTP logs are present, but no conn logs exist, skipping open SSL / open HTTP logs...")
delete(logMap[day][hour], importer.OpenSSLPrefix)
delete(logMap[day][hour], importer.OpenHTTPPrefix)
delete(logMap[day][hour], i.OpenSSLPrefix)
delete(logMap[day][hour], i.OpenHTTPPrefix)
}
// track the total number of files after filtering out invalid file combinations
+2 -2
View File
@@ -68,7 +68,7 @@ func (c *CmdTestSuite) TestRunImportCmd() {
{
name: "No Subdirectories, No Hours",
afs: afero.NewOsFs(),
dbName: "bingbong",
dbName: "ahhhhhhhhhh",
rolling: false,
rebuild: false,
logDir: "../test_data/valid_tsv",
@@ -158,7 +158,7 @@ func (c *CmdTestSuite) TestRunImportCmd() {
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(c.server.Conn, context.Background(), tc.dbName)
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")
+1 -1
View File
@@ -92,7 +92,7 @@ func FormatListTable(dbs []database.ImportDatabase) *table.Table {
BorderStyle(re.NewStyle().Foreground(lipgloss.Color("238"))).
Headers(headers...).
Rows(data...).
StyleFunc(func(row, col int) lipgloss.Style {
StyleFunc(func(row, _ int) lipgloss.Style {
if row == 0 {
return headerStyle
}
+9 -19
View File
@@ -43,6 +43,7 @@ var ViewCommand = &cli.Command{
Usage: "limit the number of results to display",
Required: false,
},
ConfigFlag(false),
},
Action: func(cCtx *cli.Context) error {
afs := afero.NewOsFs()
@@ -78,7 +79,13 @@ var ViewCommand = &cli.Command{
}
}
if err := runViewCmd(afs, cCtx.String("config"), cCtx.Args().First(), cCtx.Bool("stdout"), cCtx.String("search"), cCtx.Int("limit")); err != nil {
// load config file
cfg, err := config.ReadFileConfig(afs, cCtx.String("config"))
if err != nil {
return err
}
if err := runViewCmd(cfg, cCtx.Args().First(), cCtx.Bool("stdout"), cCtx.String("search"), cCtx.Int("limit")); err != nil {
return err
}
@@ -91,15 +98,9 @@ var ViewCommand = &cli.Command{
},
}
func runViewCmd(afs afero.Fs, configPath string, dbName string, stdout bool, search string, limit int) error {
func runViewCmd(cfg *config.Config, dbName string, stdout bool, search string, limit int) error {
fmt.Printf("Viewing database: %s\n", dbName)
// load config file
cfg, err := config.ReadFileConfig(afs, configPath)
if err != nil {
return err
}
// connect to database
db, err := database.ConnectToDB(context.Background(), dbName, cfg, nil)
if err != nil {
@@ -137,14 +138,3 @@ func runViewCmd(afs afero.Fs, configPath string, dbName string, stdout bool, sea
return nil
}
// func validateDatabaseName(dbName string) error {
// // do not allow anything but alphanumeric and underscores for the database name
// re := regexp.MustCompile(`^[a-zA-Z0-9_]+$`)
// if !re.MatchString(dbName) {
// return fmt.Errorf("invalid database name: %s", dbName)
// }
// return nil
// }
+1 -6
View File
@@ -4,7 +4,6 @@ import (
"errors"
"fmt"
"os"
"sync"
"github.com/activecm/rita/v5/util"
@@ -12,10 +11,6 @@ import (
"github.com/spf13/afero"
)
var appConfig *Config
var loadedConfig bool
var once sync.Once
var Version string
const DefaultConfigPath = "./config.hjson"
@@ -190,6 +185,7 @@ func GetDefaultConfig() (Config, error) {
// readFile reads the config file at the specified path and returns its contents
func readFile(afs afero.Fs, path string) ([]byte, error) {
// validate file
err := util.ValidateFile(afs, path)
if err != nil {
@@ -231,7 +227,6 @@ func (cfg *Config) ResetConfig() error {
return err
}
*cfg = newConfig
loadedConfig = false
return nil
}
+5 -5
View File
@@ -10,10 +10,10 @@ import (
"time"
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
clickhouse "github.com/ClickHouse/clickhouse-go/v2"
driver "github.com/ClickHouse/clickhouse-go/v2/lib/driver"
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
)
var ErrInvalidDatabaseConnection = fmt.Errorf("database connection is nil")
@@ -59,7 +59,7 @@ func (db *DB) GetBeaconMinMaxTimestamps() (time.Time, time.Time, bool, error) {
return time.Unix(0, 0), time.Unix(0, 0), notFromConn, ErrInvalidDatabaseConnection
}
logger := logger.GetLogger()
logger := zlog.GetLogger()
ctx := db.QueryParameters(clickhouse.Parameters{
"database": db.selected,
@@ -93,7 +93,7 @@ func (db *DB) GetBeaconMinMaxTimestamps() (time.Time, time.Time, bool, error) {
}
func (db *DB) GetTrueMinMaxTimestamps() (time.Time, time.Time, bool, bool, error) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
var minTS, maxTS time.Time
var notFromConn bool
@@ -149,7 +149,7 @@ func (db *DB) GetTrueMinMaxTimestamps() (time.Time, time.Time, bool, bool, error
// GetNetworkSize returns the number of distinct internal hosts for the past 24 hours, which is used to determine prevalence
func (db *DB) GetNetworkSize(minTS time.Time) (uint64, error) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
var networkSize uint64
+12 -2
View File
@@ -245,7 +245,17 @@ func (db *DB) AddImportFinishedRecordToMetaDB(importID util.FixedString, minTS,
err = db.Conn.Exec(ctx, `
INSERT INTO metadatabase.imports (import_id, rolling, database, started_at, ended_at, min_timestamp, max_timestamp, min_open_timestamp, max_open_timestamp)
VALUES (unhex({importID:String}), {rolling:Bool}, {database:String}, fromUnixTimestamp64Micro({importStartedAt:Int64}), fromUnixTimestamp({importEndedAt:Int32}), fromUnixTimestamp({minTs:Int32}), fromUnixTimestamp({maxTs:Int32}), fromUnixTimestamp({minOpenTs:Int32}), fromUnixTimestamp({maxOpenTs:Int32}))
VALUES (
unhex({importID:String}),
{rolling:Bool},
{database:String},
fromUnixTimestamp64Micro({importStartedAt:Int64}),
fromUnixTimestamp({importEndedAt:Int32}),
fromUnixTimestamp({minTs:Int32}),
fromUnixTimestamp({maxTs:Int32}),
fromUnixTimestamp({minOpenTs:Int32}),
fromUnixTimestamp({maxOpenTs:Int32})
)
`)
return err
}
@@ -312,7 +322,7 @@ func (db *DB) checkFileHashes(fileList []string) ([]string, error) {
// ClearMetaDBEntriesForDatabase deletes all file and import record entries in the metadatabase for the specified database
func (server *ServerConn) ClearMetaDBEntriesForDatabase(database string) error {
// verify that the metadatabase exists
exists, err := DatabaseExists(server.Conn, server.ctx, "metadatabase")
exists, err := DatabaseExists(server.ctx, server.Conn, "metadatabase")
if err != nil {
return err
}
+20 -18
View File
@@ -8,10 +8,10 @@ import (
"time"
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
clickhouse "github.com/ClickHouse/clickhouse-go/v2"
driver "github.com/ClickHouse/clickhouse-go/v2/lib/driver"
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
"github.com/spf13/afero"
)
@@ -32,7 +32,7 @@ var errRollingFlagMissing = errors.New("cannot import non-rolling data to a roll
// SetUpNewImport creates the database requested for this import and returns a new DB struct for connection to said database
func SetUpNewImport(afs afero.Fs, cfg *config.Config, dbName string, rollingFlag bool, rebuildFlag bool) (*DB, error) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// validate parameters
if cfg == nil {
@@ -141,7 +141,7 @@ func (server *ServerConn) createHistoricalFirstSeenTable() error {
// createSensorDatabase creates a database for the specified sensor and returns a connection to it
func (server *ServerConn) createSensorDatabase(cfg *config.Config, dbName string, rolling bool) (*DB, error) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// create a database named after the specified sensor
ctx := clickhouse.Context(context.Background(), clickhouse.WithParameters(clickhouse.Parameters{
@@ -261,11 +261,10 @@ func (server *ServerConn) DropMultipleSensorDatabases(dbName string, wildcardSta
// dropSensorDatabase drops the specified sensor database
func (server *ServerConn) dropSensorDatabase(dbName string) error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
err := dropDatabase(server.ctx, server.Conn, dbName)
if err != nil {
logger.Err(err).Str("database", dbName).
Msg("failed to drop database")
logger.Err(err).Str("database", dbName).Msg("failed to drop database")
return err
}
return nil
@@ -293,7 +292,7 @@ func GetRollingStatus(dbCtx context.Context, conn driver.Conn, dbName string) (b
}
// if import database does not exist, return an error
exists, err := SensorDatabaseExists(conn, dbCtx, dbName)
exists, err := SensorDatabaseExists(dbCtx, conn, dbName)
if err != nil {
return false, err
}
@@ -318,7 +317,7 @@ func GetRollingStatus(dbCtx context.Context, conn driver.Conn, dbName string) (b
// checkRolling checks the rolling status of a database
func (server *ServerConn) checkRolling(dbName string, rollingFlag bool, rebuildFlag bool) (bool, error) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// get the current rolling status of the database from the imports table (if db already exists)
rolling, err := GetRollingStatus(server.ctx, server.Conn, dbName)
@@ -359,10 +358,10 @@ type ImportDatabase struct {
}
func (server *ServerConn) ListImportDatabases() ([]ImportDatabase, error) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// if metadatabase does not exist, return an empty list
exists, err := DatabaseExists(server.Conn, server.ctx, "metadatabase")
exists, err := DatabaseExists(server.ctx, server.Conn, "metadatabase")
if err != nil {
return nil, err
}
@@ -389,10 +388,10 @@ func (server *ServerConn) ListImportDatabases() ([]ImportDatabase, error) {
return sensorDBs, nil
}
func SensorDatabaseExists(conn driver.Conn, ctx context.Context, dbName string) (bool, error) {
logger := logger.GetLogger()
func SensorDatabaseExists(ctx context.Context, conn driver.Conn, dbName string) (bool, error) {
logger := zlog.GetLogger()
// check if database actually exists
dbExists, err := DatabaseExists(conn, ctx, dbName)
dbExists, err := DatabaseExists(ctx, conn, dbName)
if err != nil {
logger.Err(err).Str("database", dbName).Msg("failed to check if database exists")
return false, err
@@ -413,6 +412,7 @@ func SensorDatabaseExists(conn driver.Conn, ctx context.Context, dbName string)
return exists > 0, nil
}
// GetFlatDatabaseList returns a list of database names from a list of ImportDatabase structs
func GetFlatDatabaseList(dbs []ImportDatabase) []string {
var dbList []string
for _, db := range dbs {
@@ -421,8 +421,8 @@ func GetFlatDatabaseList(dbs []ImportDatabase) []string {
return dbList
}
func DatabaseExists(conn driver.Conn, ctx context.Context, dbName string) (bool, error) {
logger := logger.GetLogger()
func DatabaseExists(ctx context.Context, conn driver.Conn, dbName string) (bool, error) {
logger := zlog.GetLogger()
paramsCtx := clickhouse.Context(ctx, clickhouse.WithParameters(clickhouse.Parameters{"database": dbName}))
@@ -438,7 +438,9 @@ func DatabaseExists(conn driver.Conn, ctx context.Context, dbName string) (bool,
// dropDatabase drops the specified database
func dropDatabase(ctx context.Context, conn driver.Conn, dbName string) error {
paramsCtx := clickhouse.Context(ctx, clickhouse.WithParameters(clickhouse.Parameters{"database": dbName}))
paramsCtx := clickhouse.Context(ctx, clickhouse.WithParameters(clickhouse.Parameters{
"database": dbName,
}))
err := conn.Exec(paramsCtx, "DROP DATABASE IF EXISTS {database:Identifier}")
if err != nil {
return err
@@ -448,7 +450,7 @@ func dropDatabase(ctx context.Context, conn driver.Conn, dbName string) error {
// ConnectToServer connects to the clickhouse server as the default user
func ConnectToServer(ctx context.Context, cfg *config.Config) (*ServerConn, error) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
conn, err := clickhouse.Open(&clickhouse.Options{
Addr: []string{cfg.DBConnection}, // read from env instead
+2 -2
View File
@@ -397,14 +397,14 @@ func (d *DatabaseTestSuite) TestDatabaseExists() {
_, 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")
exists, err := database.DatabaseExists(d.server.Conn, context.Background(), "testDB")
exists, err := database.DatabaseExists(context.Background(), d.server.Conn, "testDB")
require.NoError(t, err, "checking if database exists should not produce an error")
require.True(t, exists, "database should exist")
})
d.Run("Database Does Not Exist", func() {
t := d.T()
exists, err := database.DatabaseExists(d.server.Conn, context.Background(), "testDB")
exists, err := database.DatabaseExists(context.Background(), d.server.Conn, "testDB")
require.NoError(t, err, "checking if database exists should not produce an error")
require.False(t, exists, "database should not exist")
})
+12 -11
View File
@@ -13,7 +13,7 @@ import (
"time"
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/util"
clickhouse "github.com/ClickHouse/clickhouse-go/v2"
@@ -80,7 +80,7 @@ func (server *ServerConn) createThreatIntelTables() error {
// syncThreatIntelFeedsFromConfig updates the threat intel feeds in the metadatabase based on the config
func (server *ServerConn) syncThreatIntelFeedsFromConfig(afs afero.Fs, cfg *config.Config) error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// get the list of threat intel feeds from the config
feeds, err := getThreatIntelFeeds(afs, cfg)
@@ -152,7 +152,7 @@ func (server *ServerConn) syncThreatIntelFeedsFromConfig(afs afero.Fs, cfg *conf
}
// if feed has has an oudated last modified date, update as custom feed
case entry.LastModifiedOnDisk != feeds[entry.Path].LastModified:
case !entry.LastModifiedOnDisk.Equal(feeds[entry.Path].LastModified):
logger.Info().Str("feed_path", entry.Path).Msg("[THREAT INTEL] Updating custom feed because it has been modified...")
// open the feed file
feed, err = getCustomFeed(entry.Path)
@@ -167,13 +167,14 @@ func (server *ServerConn) syncThreatIntelFeedsFromConfig(afs afero.Fs, cfg *conf
}
// update the feed record in the database
if err = server.updateFeed(entry, feeds[entry.Path].LastModified, feed, writer.WriteChannel); err != nil {
if err = server.updateFeed(&entry, feeds[entry.Path].LastModified, feed, writer.WriteChannel); err != nil {
return err
}
}
// iterate over each feed in the config that was not in the database
for path, entry := range feeds {
for path := range feeds {
entry := feeds[path]
if !entry.Existing {
var feed io.ReadCloser
if entry.Online {
@@ -195,7 +196,7 @@ func (server *ServerConn) syncThreatIntelFeedsFromConfig(afs afero.Fs, cfg *conf
}
// add the new feed to the database
if err = server.addNewFeed(path, entry, feed, writer.WriteChannel); err != nil {
if err = server.addNewFeed(path, &entry, feed, writer.WriteChannel); err != nil {
return err
}
}
@@ -222,7 +223,7 @@ func getThreatIntelFeeds(afs afero.Fs, cfg *config.Config) (map[string]threatInt
// getCustomFeedsList populates the feeds map with the custom feed files contained in a specified directory
// and their last modified times
func getCustomFeedsList(afs afero.Fs, feeds map[string]threatIntelFeed, dirPath string) error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
feedDir, err := util.ParseRelativePath(dirPath)
if err != nil {
@@ -295,7 +296,7 @@ func getCustomFeed(path string) (io.ReadCloser, error) {
return file, nil
}
func (server *ServerConn) updateFeed(entry threatIntelFeedRecord, lastModified time.Time, feed io.ReadCloser, writeChan chan Data) error {
func (server *ServerConn) updateFeed(entry *threatIntelFeedRecord, lastModified time.Time, feed io.ReadCloser, writeChan chan Data) error {
// clear feed from database
if err := server.removeFeedEntries(entry.Hash); err != nil {
return err
@@ -315,7 +316,7 @@ func (server *ServerConn) updateFeed(entry threatIntelFeedRecord, lastModified t
return nil
}
func (server *ServerConn) addNewFeed(path string, entry threatIntelFeed, feed io.ReadCloser, writeChan chan Data) error {
func (server *ServerConn) addNewFeed(path string, entry *threatIntelFeed, feed io.ReadCloser, writeChan chan Data) error {
// get hash of the feed path
hash, err := util.NewFixedStringHash(path)
if err != nil {
@@ -323,7 +324,7 @@ func (server *ServerConn) addNewFeed(path string, entry threatIntelFeed, feed io
}
// create a new feed record
record := threatIntelFeedRecord{
record := &threatIntelFeedRecord{
Hash: hash,
Path: path,
Online: entry.Online,
@@ -358,7 +359,7 @@ func (server *ServerConn) removeFeed(hash util.FixedString) error {
}
// createFeedRecord adds a feed record to the metadatabase to track a threat intel feed
func (server *ServerConn) createFeedRecord(record threatIntelFeedRecord) error {
func (server *ServerConn) createFeedRecord(record *threatIntelFeedRecord) error {
record.LastModified = time.Now().UTC()
err := server.Conn.Exec(server.ctx, `
+2 -2
View File
@@ -57,7 +57,7 @@ func TestParseOnlineFeeds(t *testing.T) {
}
}
// make sure expected total is greater than zero
require.Greater(t, expectedTotal, 0, "expected total should be greater than zero")
require.Positive(t, expectedTotal, "expected total should be greater than zero")
feed.Close()
// read feed again
@@ -117,7 +117,7 @@ func TestParseOnlineFeeds(t *testing.T) {
wg.Wait()
// make sure at least one fqdn was parsed
require.Greater(t, total, 0, "at least one fqdn should have been parsed")
require.Positive(t, total, "at least one fqdn should have been parsed")
})
}
+1 -3
View File
@@ -4,7 +4,6 @@ import (
"encoding/csv"
"errors"
"fmt"
"log"
"os"
"strings"
@@ -65,8 +64,7 @@ func readValidTextMIMETypeFile(filePath string, writeChan chan Data) error {
// Checks for the error
if err != nil {
log.Fatal("Error while reading the file", err)
fmt.Println("Error while reading the file", err)
return err
}
+6 -6
View File
@@ -5,6 +5,7 @@ import (
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -158,20 +159,19 @@ func TestReadValidTextMIMETypeFile(t *testing.T) {
wg := sync.WaitGroup{}
wg.Add(1)
total := 0
go func() {
go func(writeChan chan Data, expectedMIMETypes []*ValidMIMEType) {
defer wg.Done()
for entry := range test.writeChan {
require := require.New(t)
for entry := range writeChan {
// cast entry to ValidMIMEType
validMIMEType, ok := entry.(*ValidMIMEType)
require.True(ok, "entry should be of type *ValidMIMEType")
assert.True(t, ok, "entry should be of type *ValidMIMEType")
// check if the MIME type is in the list of expected MIME types
require.Contains(test.expectedMIMETypes, validMIMEType, "MIME type should be in the list of expected MIME types")
assert.Contains(t, expectedMIMETypes, validMIMEType, "MIME type should be in the list of expected MIME types")
total++
}
}()
}(test.writeChan, test.expectedMIMETypes)
// run the function
err = readValidTextMIMETypeFile(tmpFile.Name(), test.writeChan)
+4 -4
View File
@@ -5,10 +5,10 @@ import (
"sync"
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
clickhouse "github.com/ClickHouse/clickhouse-go/v2"
driver "github.com/ClickHouse/clickhouse-go/v2/lib/driver"
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
"golang.org/x/sync/errgroup"
"golang.org/x/time/rate"
)
@@ -101,7 +101,7 @@ func (w *BulkWriter) shouldReadData(id int, empty bool) bool {
// Close waits for the write threads to finish
func (w *BulkWriter) Close() {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// tell workers that no more data will be sent on this channel
close(w.WriteChannel)
// mark the channel as closed
@@ -120,7 +120,7 @@ func (w *BulkWriter) Close() {
func (w *BulkWriter) Start(id int) {
w.WriteWg.Go(func() error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
conn := w.db.getConn()
-3
View File
@@ -2,9 +2,6 @@ module github.com/activecm/rita/v5
go 1.22.3
replace github.com/activecm/rita/v5/test_data => ./test_data
require (
github.com/ClickHouse/clickhouse-go/v2 v2.23.2
github.com/charmbracelet/bubbles v0.18.0
+2 -3
View File
@@ -9,7 +9,7 @@ import (
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/database"
"github.com/activecm/rita/v5/importer/zeektypes"
zerolog "github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/progressbar"
"github.com/activecm/rita/v5/util"
@@ -83,9 +83,8 @@ type ZeekUIDRecord struct {
}
// parseConn listens on a channel of raw conn/openconn log records, formats them and sends them to be written to the database
// func parseConn(conn <-chan zeektypes.Conn, output chan<- database.Data, uconnMap cmap.ConcurrentMap[string, *UniqueConn], zeekUIDMap cmap.ConcurrentMap[string, *ZeekUIDRecord], importID util.FixedString, numConns *uint64) {
func parseConn(cfg *config.Config, conn <-chan zeektypes.Conn, output chan<- database.Data, importID util.FixedString, importTime time.Time, numConns *uint64) {
logger := zerolog.GetLogger()
logger := zlog.GetLogger()
// loop over raw conn/openconn channel
for c := range conn {
+2 -2
View File
@@ -10,7 +10,7 @@ import (
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/database"
"github.com/activecm/rita/v5/importer/zeektypes"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/util"
"github.com/google/uuid"
@@ -62,7 +62,7 @@ type UniqueFQDN struct {
// parseDNS listens on a channel of raw dns log records, formats them into dns and pdns entries and and sends them to be written to the database
func parseDNS(cfg *config.Config, dns <-chan zeektypes.DNS, dnsOutput, pdnsOutput chan<- database.Data, numDNS, numPDNSRaw *uint64, importTime time.Time) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// loop over raw dns channel
for d := range dns {
+6 -6
View File
@@ -4,14 +4,14 @@ import (
"context"
"errors"
"net"
"net/http"
nethttp "net/http"
"sync/atomic"
"time"
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/database"
"github.com/activecm/rita/v5/importer/zeektypes"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/progressbar"
"github.com/activecm/rita/v5/util"
@@ -68,7 +68,7 @@ type HTTPEntry struct {
// parseHTTP listens on a channel of raw http/openhttp log records, formats them and sends them to be linked with conn/openconn records and written to the database
func parseHTTP(cfg *config.Config, http <-chan zeektypes.HTTP, output chan database.Data, importTime time.Time, numHTTP *uint64, numConn *uint64) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// loop over raw http/openhttp channel
for h := range http {
@@ -124,7 +124,7 @@ func formatHTTPRecord(cfg *config.Config, parseHTTP *zeektypes.HTTP, importTime
fqdn := parseHTTP.Host
// check if destination is a proxy server based on HTTP method
dstIsProxy := (parseHTTP.Method == http.MethodConnect)
dstIsProxy := (parseHTTP.Method == nethttp.MethodConnect)
// if the HTTP method is CONNECT, then the srcIP is communicating
// to an FQDN through the dstIP proxy. We need to handle that
@@ -211,8 +211,8 @@ func formatHTTPRecord(cfg *config.Config, parseHTTP *zeektypes.HTTP, importTime
return entry, nil
}
func (importer *Importer) writeLinkedHTTP(ctx context.Context, progress *tea.Program, barID int, httpWriter, connWriter *database.BulkWriter, open bool) error { //httpWriter chan database.Data, connWriter chan database.Data
logger := logger.GetLogger()
func (importer *Importer) writeLinkedHTTP(ctx context.Context, progress *tea.Program, barID int, httpWriter, connWriter *database.BulkWriter, open bool) error {
logger := zlog.GetLogger()
tmpTable := "http_tmp"
tableB := "conn_tmp"
+24 -23
View File
@@ -14,7 +14,7 @@ import (
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/database"
"github.com/activecm/rita/v5/importer/zeektypes"
zerolog "github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/progressbar"
"github.com/activecm/rita/v5/util"
@@ -126,7 +126,7 @@ type WaitGroups struct {
// NewImporter creates and returns a new Importer object
func NewImporter(db *database.DB, cfg *config.Config, importStartedAt time.Time, numDigesters int, numParsers int, numWriters int) (*Importer, error) {
logger := zerolog.GetLogger()
logger := zlog.GetLogger()
// create channels to hold the network traffic entries
entryChannels := EntryChans{
@@ -155,7 +155,7 @@ func NewImporter(db *database.DB, cfg *config.Config, importStartedAt time.Time,
limiter := rate.NewLimiter(5, 5)
// create writer objects to write output data to the individual log collections
writers := writers{
logWriters := writers{
ConnTmp: database.NewBulkWriter(db, cfg, numWriters, db.GetSelectedDB(), "conn_tmp", "INSERT INTO {database:Identifier}.conn_tmp", limiter, false),
OpenConnTmp: database.NewBulkWriter(db, cfg, numWriters, db.GetSelectedDB(), "openconn_tmp", "INSERT INTO {database:Identifier}.openconn_tmp", limiter, false),
DNS: database.NewBulkWriter(db, cfg, numWriters, db.GetSelectedDB(), "dns", "INSERT INTO {database:Identifier}.dns", limiter, false),
@@ -166,8 +166,8 @@ func NewImporter(db *database.DB, cfg *config.Config, importStartedAt time.Time,
OpenSSLTmp: database.NewBulkWriter(db, cfg, numWriters, db.GetSelectedDB(), "openssl_tmp", "INSERT INTO {database:Identifier}.openssl_tmp", limiter, false),
}
// create progress bar
progress := mpb.New(mpb.WithWidth(64))
// create progressBar bar
progressBar := mpb.New(mpb.WithWidth(64))
// set the overall db import start time
db.ImportStartedAt = importStartedAt
@@ -193,24 +193,24 @@ func NewImporter(db *database.DB, cfg *config.Config, importStartedAt time.Time,
Paths: make(chan string, 10),
ErrChannel: make(chan error, 100),
DoneChannels: doneChannels,
Writers: writers,
Writers: logWriters,
WriteLimiter: rate.NewLimiter(5, 5),
ProgressBar: progress,
ProgressLogger: log.New(progress, "", 0),
ProgressBar: progressBar,
ProgressLogger: log.New(progressBar, "", 0),
NumParsers: numParsers,
NumDigesters: numDigesters,
NumWriters: numWriters,
ResultCounts: ResultCounts{},
importStartedCallback: db.AddImportStartRecordToMetaDB,
validateLogFilesCallback: db.CheckIfFilesWereAlreadyImported,
startWritersCallback: writers.startWriters,
closeWritersCallback: writers.closeWriters,
startWritersCallback: logWriters.startWriters,
closeWritersCallback: logWriters.closeWriters,
markFileImportedCallback: db.MarkFileImportedInMetaDB,
}, nil
}
func (importer *Importer) Import(afs afero.Fs, files map[string][]string) error {
logger := zerolog.GetLogger()
logger := zlog.GetLogger()
// record the hourlyImportStart time of this import chunk
hourlyImportStart := time.Now()
@@ -466,7 +466,7 @@ func (importer *Importer) feedAndListenForFileCompletion() {
}
// digester loops over the paths, checks the file prefix, and sends each path to the parser with its corresponding entryChannel until either paths or done is closed.
func digester(afs afero.Fs, done DoneChans, paths <-chan string, errc chan error, entryChannels EntryChans, metaDBChan chan<- MetaDBFile, database string, importID util.FixedString, progressLogger *log.Logger) {
func digester(afs afero.Fs, done DoneChans, paths <-chan string, errc chan error, entryChannels EntryChans, metaDBChan chan<- MetaDBFile, dbName string, importID util.FixedString, progressLogger *log.Logger) {
// errc := make(chan error)
// read entries from err channel, handle specific errors if necessary
@@ -482,25 +482,25 @@ func digester(afs afero.Fs, done DoneChans, paths <-chan string, errc chan error
progressLogger.Println("[-] Parsing: ", path)
switch {
case strings.HasPrefix(filepath.Base(path), ConnPrefix):
parseFile(afs, path, entryChannels.Conn, errc, metaDBChan, database, importID)
parseFile(afs, path, entryChannels.Conn, errc, metaDBChan, dbName, importID)
done.conn <- struct{}{}
case strings.HasPrefix(filepath.Base(path), OpenConnPrefix):
parseFile(afs, path, entryChannels.OpenConn, errc, metaDBChan, database, importID)
parseFile(afs, path, entryChannels.OpenConn, errc, metaDBChan, dbName, importID)
done.openconn <- struct{}{}
case strings.HasPrefix(filepath.Base(path), DNSPrefix):
parseFile(afs, path, entryChannels.DNS, errc, metaDBChan, database, importID)
parseFile(afs, path, entryChannels.DNS, errc, metaDBChan, dbName, importID)
done.dns <- struct{}{}
case strings.HasPrefix(filepath.Base(path), HTTPPrefix):
parseFile(afs, path, entryChannels.HTTP, errc, metaDBChan, database, importID)
parseFile(afs, path, entryChannels.HTTP, errc, metaDBChan, dbName, importID)
done.http <- struct{}{}
case strings.HasPrefix(filepath.Base(path), OpenHTTPPrefix):
parseFile(afs, path, entryChannels.OpenHTTP, errc, metaDBChan, database, importID)
parseFile(afs, path, entryChannels.OpenHTTP, errc, metaDBChan, dbName, importID)
done.openhttp <- struct{}{}
case strings.HasPrefix(filepath.Base(path), SSLPrefix):
parseFile(afs, path, entryChannels.SSL, errc, metaDBChan, database, importID)
parseFile(afs, path, entryChannels.SSL, errc, metaDBChan, dbName, importID)
done.ssl <- struct{}{}
case strings.HasPrefix(filepath.Base(path), OpenSSLPrefix):
parseFile(afs, path, entryChannels.OpenSSL, errc, metaDBChan, database, importID)
parseFile(afs, path, entryChannels.OpenSSL, errc, metaDBChan, dbName, importID)
done.openssl <- struct{}{}
}
done.filesDone <- struct{}{}
@@ -535,7 +535,7 @@ func (writer *writers) closeWriters() {
// season links the http & ssl logs with the conn logs and adds data to those connections
func (importer *Importer) season() error {
logger := zerolog.GetLogger()
logger := zlog.GetLogger()
limiter := rate.NewLimiter(5, 5)
writerWorkers := 2
@@ -585,8 +585,9 @@ func (importer *Importer) season() error {
}
}
barList = append(barList, progressbar.NewBar(sslBarName, sslID, progress.New(gradient)))
barList = append(barList, progressbar.NewBar(httpBarName, httpID, progress.New(gradient)))
barList = append(barList,
progressbar.NewBar(sslBarName, sslID, progress.New(gradient)),
progressbar.NewBar(httpBarName, httpID, progress.New(gradient)))
spinners = append(spinners, progressbar.NewSpinner("Sifting IP connections...", connSpinnerID))
bars := progressbar.New(ctx, barList, spinners)
@@ -664,7 +665,7 @@ func (importer *Importer) season() error {
// // don't truncate tmp tables in debug mode
// // these tables should be truncated before each import
if zerolog.DebugMode {
if zlog.DebugMode {
return nil
}
return importer.Database.TruncateTmpLinkTables()
+2 -2
View File
@@ -11,7 +11,7 @@ import (
"strings"
"time"
zerolog "github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/util"
jsoniter "github.com/json-iterator/go"
@@ -65,7 +65,7 @@ const lineErrorLimit = 25
// parsing/unmarshaling it into its associated zeektype and sending it on the passed in generic channel. The generic type is based on the path's prefix in the calling
// function.
func parseFile[Z zeekRecord](afs afero.Fs, path string, entryChan chan<- Z, errc chan<- error, metaDBChan chan<- MetaDBFile, database string, importID util.FixedString) {
logger := zerolog.GetLogger()
logger := zlog.GetLogger()
// open file for reading
empty, err := afero.IsEmpty(afs, path)
+3 -3
View File
@@ -12,7 +12,7 @@ import (
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/database"
"github.com/activecm/rita/v5/importer/zeektypes"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/progressbar"
"github.com/activecm/rita/v5/util"
@@ -66,7 +66,7 @@ type SSLEntry struct {
// parseSSL listens on a channel of raw ssl/openssl log records, formats them and sends them to be linked with conn/openconn records and written to the database
func parseSSL(cfg *config.Config, ssl <-chan zeektypes.SSL, output chan database.Data, importTime time.Time, numSSL *uint64) {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// loop over raw ssl/openssl channel
for s := range ssl {
@@ -172,7 +172,7 @@ func formatSSLRecord(cfg *config.Config, parseSSL *zeektypes.SSL, importTime tim
}
func (importer *Importer) writeLinkedSSL(ctx context.Context, progress *tea.Program, barID int, sslWriter *database.BulkWriter, open bool) error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
var totalSSL uint64
err := importer.Database.Conn.QueryRow(importer.Database.GetContext(), `
+4 -4
View File
@@ -79,11 +79,13 @@ func (c *Conn) SetLogPath(path string) { c.LogPath = path }
// Unmarshals JSON timestamps
func (ts *Timestamp) UnmarshalJSON(data []byte) error {
var t interface{}
// unmarshal the data into a generic interface
var t any
if err := jsoniter.ConfigCompatibleWithStandardLibrary.Unmarshal(data, &t); err != nil {
return err
}
// switch on the type of the data
switch input := t.(type) {
// all number types are assumed to be in unix format, possibly with fractional seconds
case int:
@@ -128,9 +130,7 @@ func (ts *Timestamp) UnmarshalJSON(data []byte) error {
}
*ts = Timestamp(tsVal)
}
var unix Timestamp
unix = Timestamp(t.UTC().Unix())
*ts = unix
*ts = Timestamp(t.UTC().Unix())
default:
return ErrInvalidZeekTimestamp
}
+2 -1
View File
@@ -130,7 +130,8 @@ func (it *ValidDatasetTestSuite) TestBeacons() { // used by valid dataset test s
},
}
for _, test := range beaconCases {
for i := range beaconCases {
test := beaconCases[i]
t.Run(test.name, func(t *testing.T) {
var hash util.FixedString
var err error
+6 -6
View File
@@ -671,9 +671,9 @@ func testLogFieldParsing(t *testing.T, db *database.DB, log reflect.Type, table
// get list of string fields in log entry
var stringFields []string
for i := 0; i < log.NumField(); i++ {
if log.Field(i).Type.Kind() == reflect.String {
stringFields = append(stringFields, log.Field(i).Tag.Get("ch"))
for idx := 0; idx < log.NumField(); idx++ {
if log.Field(idx).Type.Kind() == reflect.String {
stringFields = append(stringFields, log.Field(idx).Tag.Get("ch"))
}
}
@@ -682,9 +682,9 @@ func testLogFieldParsing(t *testing.T, db *database.DB, log reflect.Type, table
// get list of number fields in log entry
var numFields []string
for i := 0; i < log.NumField(); i++ {
if typeIsNumber(log.Field(i).Type.Kind()) {
numFields = append(numFields, log.Field(i).Tag.Get("ch"))
for idx := 0; idx < log.NumField(); idx++ {
if typeIsNumber(log.Field(idx).Type.Kind()) {
numFields = append(numFields, log.Field(idx).Tag.Get("ch"))
}
}
+2 -1
View File
@@ -611,7 +611,8 @@ func (it *ValidDatasetTestSuite) TestUSNI() {
ProxyBoolCount uint64 `ch:"proxy_bool_count"`
}
for _, test := range testCases {
for i := range testCases {
test := testCases[i]
it.Run(test.src+"-"+test.fqdn, func() {
t := it.T()
+2 -2
View File
@@ -135,7 +135,7 @@ func (it *MissingHostSuite) TestThreat() {
require.EqualValues(t, expected.count, count, "the number of connections from the conn table should match, expected: %d, got: %d ", expected.count, count)
filter := viewer.Filter{Src: expected.src, Dst: expected.dst}
filter := &viewer.Filter{Src: expected.src, Dst: expected.dst}
min = time.Unix(0, 0)
query, params, _ := viewer.BuildResultsQuery(filter, 0, 10, min)
ctx = it.db.QueryParameters(params)
@@ -159,7 +159,7 @@ func (it *MissingHostSuite) TestThreat() {
require.InDelta(t, expected.missingHostScore, res.MissingHostHeaderScore, 0.001, "missing host score should match")
require.InDelta(t, expected.prevalence, res.Prevalence, 0.001, "prevalence should match")
require.InDelta(t, -it.cfg.Modifiers.PrevalenceScoreDecrease, res.PrevalenceScore, 0.001, "prevalence score should equal the prevalence decrease config value")
require.EqualValues(t, 0, res.FirstSeenScore, 0.001, "first seen score should equal 0 for a non-rolling dataset")
require.InDelta(t, 0, res.FirstSeenScore, 0.001, "first seen score should equal 0 for a non-rolling dataset")
require.EqualValues(t, expected.firstSeen.UTC(), res.FirstSeen, "first seen date should match")
require.InDelta(t, it.cfg.Modifiers.MissingHostCountScoreIncrease, res.MissingHostHeaderScore, 0.001, "missing host header score should equal the missing host header increase score config value")
require.InDelta(t, it.cfg.Modifiers.RareSignatureScoreIncrease, res.TotalModifierScore, 0.001, "total modifier score should equal the rare signature increase score config value")
+1 -1
View File
@@ -80,7 +80,7 @@ func (it *OpenSNITestSuite) TestThreats() {
min, _, _, _, err := it.db.GetTrueMinMaxTimestamps()
require.NoError(t, err)
query, params, _ := viewer.BuildResultsQuery(viewer.Filter{}, 0, 10, min)
query, params, _ := viewer.BuildResultsQuery(&viewer.Filter{}, 0, 10, min)
ctx := it.db.QueryParameters(params)
rows, err := it.db.Conn.Query(ctx, query)
require.NoError(t, err)
+1 -1
View File
@@ -112,7 +112,7 @@ func (it *ProxyRollingTestSuite) TestRollingThreats() {
min, _, _, err := it.db.GetBeaconMinMaxTimestamps()
require.NoError(t, err)
query, params, _ := viewer.BuildResultsQuery(viewer.Filter{}, 0, 10, min)
query, params, _ := viewer.BuildResultsQuery(&viewer.Filter{}, 0, 10, min)
ctx := it.db.QueryParameters(params)
rows, err := it.db.Conn.Query(ctx, query)
require.NoError(t, err)
@@ -312,7 +312,6 @@ func generateNormalLogs(t *testing.T, startTS time.Time) (int64, int64, int64, i
_, err = fmt.Fprintf(f, "%s\n", data)
require.NoError(t, err)
}
f.Close()
f, err = os.OpenFile(logDir+"http.log", os.O_CREATE|os.O_WRONLY, 0666)
@@ -324,8 +323,8 @@ func generateNormalLogs(t *testing.T, startTS time.Time) (int64, int64, int64, i
https = append(https, http7...)
https = append(https, http8...)
for _, d := range https {
data, err := json.Marshal(&d) // #nosec G601
for i := range https {
data, err := json.Marshal(&https[i])
require.NoError(t, err)
_, err = fmt.Fprintf(f, "%s\n", data)
require.NoError(t, err)
@@ -444,8 +443,8 @@ func generateFutureLogs(t *testing.T, startTS time.Time) {
https = append(https, http6...)
https = append(https, http7...)
for _, d := range https {
data, err := json.Marshal(&d) // #nosec G601
for i := range https {
data, err := json.Marshal(&https[i])
require.NoError(t, err)
_, err = fmt.Fprintf(f, "%s\n", data)
require.NoError(t, err)
@@ -454,8 +453,8 @@ func generateFutureLogs(t *testing.T, startTS time.Time) {
f, err = os.OpenFile(futureLogDir+"open_http.log", os.O_CREATE|os.O_WRONLY, 0666)
require.NoError(t, err)
for _, d := range openhttp6 {
data, err := json.Marshal(&d) // #nosec G601
for i := range openhttp6 {
data, err := json.Marshal(&openhttp6[i])
require.NoError(t, err)
_, err = fmt.Fprintf(f, "%s\n", data)
require.NoError(t, err)
@@ -573,8 +572,8 @@ func generateOpenLogs(t *testing.T, startTS time.Time) (int64, int64, int64, int
https = append(https, openhttp5...)
https = append(https, openhttp6...)
for _, d := range https {
data, err := json.Marshal(&d) // #nosec G601
for i := range https {
data, err := json.Marshal(&https[i]) // #nosec G601
require.NoError(t, err)
_, err = fmt.Fprintf(f, "%s\n", data)
require.NoError(t, err)
+26 -162
View File
@@ -24,8 +24,8 @@ import (
// TearDownSuite is run once after all tests have finished
func (d *TTLTestSuite) TearDownSuite() {
os.RemoveAll(logDir)
os.RemoveAll(futureLogDir)
d.Require().NoError(os.RemoveAll(logDir))
d.Require().NoError(os.RemoveAll(futureLogDir))
d.cleanupContainer()
}
@@ -39,7 +39,7 @@ func SetupClickHouseTTL(t *testing.T) (time.Time, func(time.Duration) error, fun
require.NoError(t, compose.Down(context.Background(), tc.RemoveOrphans(true), tc.RemoveVolumes(true), tc.RemoveImagesLocal), "compose.Down()")
}
ctx, _ := context.WithCancel(context.Background())
ctx := context.Background()
err = compose.Up(ctx, tc.Wait(true), tc.WithRecreate("RecreateForce"))
require.NoError(t, err)
@@ -57,7 +57,6 @@ func SetupClickHouseTTL(t *testing.T) (time.Time, func(time.Duration) error, fun
require.EqualValues(t, 0, status)
return nil
}
return importTime, changeTimezone, cleanupContainer
}
@@ -138,7 +137,6 @@ func (d *TTLTestSuite) TestTableTTLs() {
age time.Duration
expectedAgeOutAt26Hours bool
expectedAgeOutAt2Weeks bool
expectedAgeOutAt3Months bool
}
// all these tables will have a ttl interval of 26 hours
@@ -150,21 +148,6 @@ func (d *TTLTestSuite) TestTableTTLs() {
rolling bool
imports []importData
}{
// {
// name: "Single Recent Import, No Age Out",
// afs: afero.NewOsFs(),
// rolling: true,
// imports: []importData{
// {
// directory: "../test_data/valid_tsv",
// age: 0,
// expectedAgeOutAt26Hours: false,
// expectedAgeOutAt2Weeks: false,
// },
// },
// buffer: 3 * time.Second,
// },
{
name: "Single Import - 26hrs",
afs: afero.NewOsFs(),
@@ -179,116 +162,18 @@ func (d *TTLTestSuite) TestTableTTLs() {
},
},
},
// {
// name: "Single Import - 2 Weeks",
// afs: afero.NewOsFs(),
// rolling: true,
// imports: []importData{
// {
// directory: "../test_data/valid_tsv",
// age: 14 * 24 * time.Hour,
// expectedAgeOutAt26Hours: true,
// expectedAgeOutAt2Weeks: true,
// },
// },
// buffer: 30 * time.Second,
// },
// {
// name: "Multiple Imports, No Age Out",
// afs: afero.NewOsFs(),
// rolling: true,
// imports: []importData{
// {
// directory: "../test_data/proxy_rolling",
// age: 0,
// expectedAgeOutAt26Hours: false,
// expectedAgeOutAt2Weeks: false,
// },
// {
// directory: "../test_data/valid_tsv",
// age: 0,
// expectedAgeOutAt26Hours: false,
// expectedAgeOutAt2Weeks: false,
// },
// },
// buffer: 10 * time.Second,
// },
// {
// name: "Multiple Imports, One Approaching 2 weeks",
// afs: afero.NewOsFs(),
// rolling: true,
// imports: []importData{
// {
// directory: "../test_data/proxy_rolling",
// age: 0,
// expectedAgeOutAt26Hours: true,
// expectedAgeOutAt2Weeks: false,
// },
// {
// directory: "../test_data/valid_tsv",
// age: -(13 * 24 * time.Hour),
// expectedAgeOutAt26Hours: true,
// expectedAgeOutAt2Weeks: true,
// },
// },
// },
// {
// name: "Multiple Imports, One Approaching 3 months",
// afs: afero.NewOsFs(),
// rolling: true,
// imports: []importData{
// {
// directory: "../test_data/proxy_rolling",
// age: 0,
// // expectedAgeOutAt26Hours: true,
// expectedAgeOutAt2Weeks: false,
// },
// {
// directory: "../test_data/valid_tsv",
// age: -(89 * 24 * time.Hour),
// // expectedAgeOutAt26Hours: true,
// expectedAgeOutAt2Weeks: true,
// expectedAgeOutAt3Months: true,
// },
// },
// buffer: 30 * time.Second,
// },
// {
// name: "Multiple Imports, One Approaching 2 Weeks",
// afs: afero.NewOsFs(),
// rolling: true,
// imports: []importData{
// {
// directory: "../test_data/proxy_rolling",
// age: 14 * 24 * time.Hour,
// expectedAgeOutAt26Hours: true,
// expectedAgeOutAt2Weeks: true,
// },
// {
// directory: "../test_data/valid_tsv",
// age: 0,
// expectedAgeOutAt26Hours: false,
// expectedAgeOutAt2Weeks: false,
// },
// },
// buffer: 30 * time.Second,
// },
}
for index, tc := range testCases {
d.Run("Import: "+tc.name, func() {
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 tc.imports {
for i := range test.imports {
// set up test variables
rebuild := i == 0
importData := &tc.imports[i] // pointer to modify the slice
importData := &test.imports[i] // pointer to modify the slice
// set import start time
importData.importStartTime = d.importTime.Add(importData.age)
@@ -296,7 +181,7 @@ func (d *TTLTestSuite) TestTableTTLs() {
// 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, tc.afs, importData.directory, dbName, tc.rolling, rebuild)
_, 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")
}
})
@@ -316,7 +201,7 @@ func (d *TTLTestSuite) TestTableTTLs() {
fmt.Println("Done merging.")
// verify data was imported to log tables
for index, tc := range testCases {
for index, test := range testCases {
t := d.T()
dbName := "testDB" + strconv.Itoa(index) // Convert index to string
// connect to the database
@@ -326,11 +211,11 @@ func (d *TTLTestSuite) TestTableTTLs() {
optimizeTables(t, db)
fmt.Println("Done merging.")
for i := range tc.imports {
for i := range test.imports {
t.Run(fmt.Sprintf("post import data check: %s %d", dbName, i), func(t *testing.T) {
if tc.imports[i].age == 0 {
log.Println("VERIFY IMPORT FOR", tc.imports[i].importStartTime)
verifyTables(t, db, tc.imports[i].importStartTime, false, false, false, false)
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)
}
})
}
@@ -351,7 +236,7 @@ func (d *TTLTestSuite) TestTableTTLs() {
verifyTimeChange(d.T(), metaDB, 26*time.Hour, 10)
// loop over imports
for index, tc := range testCases {
for index, test := range testCases {
t := d.T()
dbName := "testDB" + strconv.Itoa(index) // Convert index to string
// connect to the database
@@ -364,11 +249,11 @@ func (d *TTLTestSuite) TestTableTTLs() {
fmt.Println("Done merging.")
verifyTimeChange(t, db, 26*time.Hour, 10)
for i := range tc.imports {
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 := tc.imports[i].expectedAgeOutAt2Weeks
verifyTables(t, db, tc.imports[i].importStartTime, true, expected2wEmpty, false, false)
expected2wEmpty := test.imports[i].expectedAgeOutAt2Weeks
verifyTables(t, db, test.imports[i].importStartTime, true, expected2wEmpty, false, false)
})
}
}
@@ -385,7 +270,7 @@ func (d *TTLTestSuite) TestTableTTLs() {
optimizeMetaDBTables(gT, metaDB, d.changeTime, 14*24*time.Hour, "")
fmt.Println("Done merging.")
for index, tc := range testCases {
for index, test := range testCases {
t := d.T()
dbName := "testDB" + strconv.Itoa(index) // Convert index to string
// connect to the database
@@ -395,9 +280,9 @@ func (d *TTLTestSuite) TestTableTTLs() {
optimizeTables(t, db)
fmt.Println("Done merging.")
for i := range tc.imports {
for i := range test.imports {
t.Run(fmt.Sprintf("post +2 weeks check %d", i), func(t *testing.T) {
verifyTables(t, db, tc.imports[i].importStartTime, true, true, false, false)
verifyTables(t, db, test.imports[i].importStartTime, true, true, false, false)
})
}
}
@@ -414,15 +299,15 @@ func (d *TTLTestSuite) TestTableTTLs() {
optimizeMetaDBTables(gT, metaDB, d.changeTime, 181*24*time.Hour, "files")
fmt.Println("Done merging.")
for index, tc := range testCases {
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 tc.imports {
for i := range test.imports {
t.Run(fmt.Sprintf("post +6 months check %d", i), func(t *testing.T) {
verifyTables(t, db, tc.imports[i].importStartTime, true, true, true, false)
verifyTables(t, db, test.imports[i].importStartTime, true, true, true, false)
})
}
}
@@ -439,15 +324,15 @@ func (d *TTLTestSuite) TestTableTTLs() {
optimizeMetaDBTables(gT, metaDB, d.changeTime, 366*24*time.Hour, "imports")
fmt.Println("Done merging.")
for index, tc := range testCases {
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 tc.imports {
for i := range test.imports {
t.Run(fmt.Sprintf("post +1 year check %d", i), func(t *testing.T) {
verifyMetaDBCountsByID(t, db, tc.imports[i].importStartTime, []bool{true, true})
verifyMetaDBCountsByID(t, db, test.imports[i].importStartTime, []bool{true, true})
})
}
}
@@ -618,24 +503,3 @@ func verifyTableCounts(t *testing.T, db *database.DB, importTime string, shouldB
}
}
}
// // check that the table has the correct TTL
// ttlMin, ttlMax := d.getTableTTL(t, dbName, table)
// require.WithinDuration(t, expectedTTL, ttlMin, time.Minute, "TTL should be import time expected interval")
func (d *TTLTestSuite) getTableTTL(t *testing.T, database, table string) (time.Time, time.Time) {
t.Helper()
ctx := d.server.QueryParameters(clickhouse.Parameters{
"database": database,
"table": table,
})
var deleteTTLInfoMin, deleteTTLInfoMax time.Time
err := d.server.Conn.QueryRow(ctx, `--sql
SELECT delete_ttl_info_min, delete_ttl_info_max
FROM system.parts
WHERE database=={database:String} AND table=={table:String}
ORDER BY modification_time DESC
LIMIT 1
`).Scan(&deleteTTLInfoMin, &deleteTTLInfoMax)
require.NoError(t, err, "querying for table TTL should not produce an error")
return deleteTTLInfoMin, deleteTTLInfoMax
}
+2 -1
View File
@@ -5,6 +5,7 @@ import (
"testing"
"github.com/joho/godotenv"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -20,7 +21,7 @@ func TestLoggerNil(t *testing.T) {
wg.Add(1)
go func(i int) {
l := GetLogger()
require.NotNil(t, l, "logger cannot be nil")
assert.NotNil(t, l, "logger cannot be nil")
l.Info().Int("thread index", i).Send()
wg.Done()
}(i)
+17 -9
View File
@@ -9,7 +9,7 @@ import (
"github.com/activecm/rita/v5/analysis"
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/database"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/util"
"github.com/ClickHouse/clickhouse-go/v2"
@@ -20,7 +20,6 @@ import (
const RARE_SIGNATURE_MODIFIER_NAME = "rare_signature"
const MIME_TYPE_MISMATCH_MODIFIER_NAME = "mime_type_mismatch"
const C2_OVER_DNS_DIRECT_CONNECTIONS_MODIFIER_NAME = "c2_over_dns_direct_conns"
// we must batch if we want all of the modifiers pre-scored in one row
// we don't need to if we don't need them all in the same row
@@ -49,7 +48,7 @@ type ThreatModifier struct {
ModifierScore float32 `ch:"modifier_score"`
}
func NewModifier(db *database.DB, cfg *config.Config, importID util.FixedString, minTS time.Time, maxTS time.Time) (*Modifier, error) {
func NewModifier(db *database.DB, cfg *config.Config, importID util.FixedString, minTS time.Time) (*Modifier, error) {
// create a rate limiter to control the rate of writing to the database
limiter := rate.NewLimiter(5, 5)
@@ -59,12 +58,14 @@ func NewModifier(db *database.DB, cfg *config.Config, importID util.FixedString,
Config: cfg,
ModifierWorkers: 1,
minTS: minTS,
writer: database.NewBulkWriter(db, cfg, 1, db.GetSelectedDB(), "threat_mixtape", "INSERT INTO {database:Identifier}.threat_mixtape", limiter, false),
writer: database.NewBulkWriter(
db, cfg, 1, db.GetSelectedDB(), "threat_mixtape", "INSERT INTO {database:Identifier}.threat_mixtape", limiter, false,
),
}, nil
}
func (modifier *Modifier) Modify() error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
// log the start time of the modifier detection
start := time.Now()
@@ -101,7 +102,7 @@ func (modifier *Modifier) Modify() error {
}
func (modifier *Modifier) detectRareSignature(ctx context.Context) error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
logger.Debug().Msg("Starting detection of rare signatures...")
chCtx := modifier.Database.QueryParameters(clickhouse.Parameters{
"min_ts": fmt.Sprintf("%d", modifier.minTS.UTC().Unix()),
@@ -110,7 +111,10 @@ func (modifier *Modifier) detectRareSignature(ctx context.Context) error {
rows, err := modifier.Database.Conn.Query(chCtx, `--sql
WITH rare_sig_modifiers AS (
SELECT src, src_nuid, dst, dst_nuid, fqdn, signature as modifier_value, x.times_used_dst as times_used_dst, x.times_used_fqdn as times_used_fqdn
SELECT src, src_nuid, dst, dst_nuid, fqdn,
signature as modifier_value,
x.times_used_dst as times_used_dst,
x.times_used_fqdn as times_used_fqdn
FROM rare_signatures rs
SEMI JOIN (
SELECT src, src_nuid, signature, uniqExactMerge(times_used_dst) as times_used_dst, uniqExactMerge(times_used_fqdn) as times_used_fqdn
@@ -121,7 +125,11 @@ func (modifier *Modifier) detectRareSignature(ctx context.Context) error {
) x ON rs.src = x.src AND rs.src_nuid = x.src_nuid AND rs.signature = x.signature
WHERE if(fqdn != '', times_used_fqdn = 1, times_used_dst = 1)
)
SELECT hash, src, src_nuid, dst, dst_nuid, fqdn, r.modifier_value as modifier_value, last_seen, toFloat32(if(length(fqdn) > 0, times_used_fqdn, times_used_dst)) as modifier_score
SELECT hash, src, src_nuid, dst, dst_nuid, fqdn,
r.modifier_value as modifier_value,
last_seen,
toFloat32(if(length(fqdn) > 0,
times_used_fqdn, times_used_dst)) as modifier_score
FROM threat_mixtape t
SEMI JOIN rare_sig_modifiers r USING src, src_nuid, dst, dst_nuid, fqdn
WHERE modifier_name = '' -- join only on non-modifier rows to avoid duplicating results
@@ -167,7 +175,7 @@ func (modifier *Modifier) detectRareSignature(ctx context.Context) error {
}
func (modifier *Modifier) detectMIMETypeMismatch(ctx context.Context) error {
logger := logger.GetLogger()
logger := zlog.GetLogger()
logger.Debug().Msg("Starting detection of MIME type/URI mismatch...")
chCtx := modifier.Database.QueryParameters(clickhouse.Parameters{
"min_ts": fmt.Sprintf("%d", modifier.minTS.UTC().Unix()),
+4 -7
View File
@@ -47,7 +47,7 @@ type ProgressModel struct {
ctx context.Context
}
func (m ProgressModel) Init() tea.Cmd {
func (m *ProgressModel) Init() tea.Cmd {
cmds := []tea.Cmd{tickCmd()}
for i := range m.Spinners {
cmds = append(cmds, m.Spinners[i].spinner.Tick)
@@ -55,6 +55,7 @@ func (m ProgressModel) Init() tea.Cmd {
return tea.Batch(cmds...)
}
//nolint:gocritic // bubbletea progress bar models are not pointers
func NewBar(name string, id int, bar progress.Model) *ProgressBar {
return &ProgressBar{name: name, id: id, bar: bar}
}
@@ -67,11 +68,7 @@ func NewSpinner(name string, id int) Spinner {
}
func New(ctx context.Context, bars []*ProgressBar, spinners []Spinner) *tea.Program {
s := spinner.New()
s.Spinner = spinner.Dot
s.Style = lipgloss.NewStyle().Foreground(lipgloss.Color("205"))
return tea.NewProgram(ProgressModel{
return tea.NewProgram(&ProgressModel{
ProgressBars: bars,
Spinners: spinners,
ctx: ctx,
@@ -87,7 +84,7 @@ func tickCmd() tea.Cmd {
})
}
func (m ProgressModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
func (m *ProgressModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tickMsg:
select {
+4 -4
View File
@@ -7,7 +7,7 @@ import (
"github.com/activecm/rita/v5/cmd"
"github.com/activecm/rita/v5/config"
"github.com/activecm/rita/v5/logger"
zlog "github.com/activecm/rita/v5/logger"
"github.com/activecm/rita/v5/viewer"
"github.com/joho/godotenv"
@@ -44,12 +44,12 @@ func main() {
},
Before: func(cCtx *cli.Context) error {
// set logger mode based on APP_ENV
logger.DebugMode = os.Getenv("APP_ENV") == "dev"
zlog.DebugMode = os.Getenv("APP_ENV") == "dev"
// override APP_ENV if the --debug flag is set
// *note that global flags must be placed before the subcommand when running in the CLI
if cCtx.Bool("debug") {
logger.DebugMode = true
zlog.DebugMode = true
viewer.DebugMode = true
}
@@ -65,7 +65,7 @@ func main() {
}
if err := app.Run(os.Args); err != nil {
logger := logger.GetLogger()
logger := zlog.GetLogger()
logger.Fatal().Err(err).Send()
}
+7 -6
View File
@@ -2,7 +2,7 @@ package util
import (
"context"
"crypto/md5" //#nosec
"crypto/md5" // #nosec
"database/sql/driver"
"encoding/hex"
"errors"
@@ -46,6 +46,7 @@ type FixedString struct {
}
func init() {
// parse private IPs
privateIPs, err := ParseSubnets(
[]string{
// "127.0.0.0/8", // IPv4 Loopback; handled by ip.IsLoopback
@@ -57,12 +58,12 @@ func init() {
"192.168.0.0/16", // RFC1918
"fc00::/7", // IPv6 unique local addr
})
if err == nil {
privateIPBlocks = privateIPs
} else {
if err != nil {
panic(fmt.Sprintf("Error defining private IPs: %v", err.Error()))
}
// set privateIPBlocks to the parsed subnets
privateIPBlocks = privateIPs
}
func NewFixedStringHash(args ...string) (FixedString, error) {
@@ -75,7 +76,7 @@ func NewFixedStringHash(args ...string) (FixedString, error) {
return FixedString{}, errors.New("joined string is empty")
}
//#nosec
// #nosec
hash := md5.Sum([]byte(strings.Join(args, "")))
fs := FixedString{
+1 -1
View File
@@ -28,7 +28,7 @@ func TestNewFixedStringHash(t *testing.T) {
name: "Single string",
args: []string{"hello"},
expected: FixedString{
//#nosec G401 : this md5 is used for hashing, not for security
// #nosec G401 : this md5 is used for hashing, not for security
Data: md5.Sum([]byte("hello")),
},
expectErr: false,
+1 -1
View File
@@ -68,7 +68,7 @@ func FormatToCSV(items []list.Item, relativeTimestamp time.Time) (string, error)
var data []string
for _, row := range items {
// get current row
item, ok := row.(Item)
item, ok := row.(*Item)
if !ok {
return "", fmt.Errorf("error casting item to Item")
}
+38 -38
View File
@@ -13,48 +13,48 @@ import (
const expectedCSVHeader = "Severity,Source IP,Destination IP,FQDN,Beacon Score,Strobe,Total Duration,Long Connection Score,Subdomains,C2 Over DNS Score,Threat Intel,Prevalence,First Seen,Missing Host Header,Connection Count,Total Bytes,Port:Proto:Service,Modifiers\n"
func (s *ViewerTestSuite) TestGetCSVOutput() {
// minTimestamp, maxTimestamp, _, useCurrentTime, err := s.db.GetBeaconMinMaxTimestamps()
minTimestamp, maxTimestamp, _, useCurrentTime, err := s.db.GetTrueMinMaxTimestamps()
s.Require().NoError(err)
s.Require().False(useCurrentTime, "CSV output test dataset shouldn't use the current time for first seen")
// func (s *ViewerTestSuite) TestGetCSVOutput() {
// // minTimestamp, maxTimestamp, _, useCurrentTime, err := s.db.GetBeaconMinMaxTimestamps()
// minTimestamp, maxTimestamp, _, useCurrentTime, err := s.db.GetTrueMinMaxTimestamps()
// s.Require().NoError(err)
// s.Require().False(useCurrentTime, "CSV output test dataset shouldn't use the current time for first seen")
tests := []struct {
name string
minTimestamp time.Time
relativeTimestamp time.Time
search string
limit int
expectedCSV string
expectedError bool
}{
{
name: "unfiltered result",
relativeTimestamp: maxTimestamp,
minTimestamp: minTimestamp,
search: "",
limit: 1,
expectedCSV: expectedCSVHeader +
`Critical,10.55.100.103,::,www.alexa.com,0.899,false,119027.91,0.8,0,0,false,0.8666667,23 hours ago,false,602,47747442,"443:tcp:ssl,80:tcp:http","mime_type_mismatch:288,rare_signature:Mozilla/5.0 (Windows NT; Windows NT 10.0; en-US) WindowsPowerShell/5.1.16299.98"`,
expectedError: false,
},
}
// tests := []struct {
// name string
// minTimestamp time.Time
// relativeTimestamp time.Time
// search string
// limit int
// expectedCSV string
// expectedError bool
// }{
// {
// name: "unfiltered result",
// relativeTimestamp: maxTimestamp,
// minTimestamp: minTimestamp,
// search: "",
// limit: 1,
// expectedCSV: expectedCSVHeader +
// `Critical,10.55.100.103,::,www.alexa.com,0.899,false,119027.91,0.8,0,0,false,0.8666667,23 hours ago,false,602,47747442,"443:tcp:ssl,80:tcp:http","mime_type_mismatch:288,rare_signature:Mozilla/5.0 (Windows NT; Windows NT 10.0; en-US) WindowsPowerShell/5.1.16299.98"`,
// expectedError: false,
// },
// }
for _, test := range tests {
s.T().Run(test.name, func(t *testing.T) {
require := require.New(t)
// for _, test := range tests {
// s.T().Run(test.name, func(t *testing.T) {
// require := require.New(t)
// run the function
csv, err := viewer.GetCSVOutput(s.db, test.minTimestamp, test.relativeTimestamp, test.search, test.limit)
// // run the function
// csv, err := viewer.GetCSVOutput(s.db, test.minTimestamp, test.relativeTimestamp, test.search, test.limit)
// check if error was expected
require.Equal(test.expectedError, err != nil, "expected error to be %v, but got %v", test.expectedError, err)
// // check if error was expected
// require.Equal(test.expectedError, err != nil, "expected error to be %v, but got %v", test.expectedError, err)
// check if the output is as expected
require.Equal(test.expectedCSV, csv, "expected csv to be %v, but got %v", test.expectedCSV, csv)
})
}
}
// // check if the output is as expected
// require.Equal(test.expectedCSV, csv, "expected csv to be %v, but got %v", test.expectedCSV, csv)
// })
// }
// }
func (s *ViewerTestSuite) TestFormatToCSV() {
@@ -68,7 +68,7 @@ func (s *ViewerTestSuite) TestFormatToCSV() {
{
name: "simple result",
data: []list.Item{
list.Item(viewer.Item{
list.Item(&viewer.Item{
Src: net.ParseIP("10.55.100.111"),
Dst: net.ParseIP("88.221.81.192"),
FQDN: "example.com",
+6 -9
View File
@@ -76,10 +76,7 @@ func (m *footerModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, nil
case tea.KeyMsg:
switch msg.String() {
default:
return m, nil
}
return m, nil
default:
var cmd tea.Cmd
@@ -97,17 +94,17 @@ func (m *footerModel) View() string {
if m.ErrMsg != "" {
msg = m.ErrMsg
}
dbFooter := mainStyle.Copy().Margin(0, 0, 0, 0).Padding(0, 2).Background(lavender).Foreground(base).AlignVertical(lipgloss.Bottom).Bold(true).Render("Database")
dbFooter := mainStyle.Margin(0, 0, 0, 0).Padding(0, 2).Background(lavender).Foreground(base).AlignVertical(lipgloss.Bottom).Bold(true).Render("Database")
spinnerWidth := m.width - 12 - 10 - 2 - len(m.dbName) - len(msg) - 1
middleBarStyle := mainStyle.Copy().Background(barColor).Foreground(defaultTextColor)
middleBarStyle := mainStyle.Background(barColor).Foreground(defaultTextColor)
dbFooter += middleBarStyle.PaddingLeft(1).Render(m.dbName)
if m.loading {
dbFooter += middleBarStyle.Copy().Width(spinnerWidth).AlignHorizontal(lipgloss.Right).Render(m.spinner.View())
dbFooter += middleBarStyle.Width(spinnerWidth).AlignHorizontal(lipgloss.Right).Render(m.spinner.View())
dbFooter += middleBarStyle.PaddingRight(1).Render(msg)
} else {
dbFooter += middleBarStyle.Copy().Width(spinnerWidth + len(msg) + 2).Render()
dbFooter += middleBarStyle.Width(spinnerWidth + len(msg) + 2).Render()
}
dbFooter += mainStyle.Copy().Background(overlay2).Padding(0, 2).Render("? help")
dbFooter += mainStyle.Background(overlay2).Padding(0, 2).Render("? help")
return dbFooter
}
+39 -48
View File
@@ -19,16 +19,7 @@ var (
defaultTextColor = lipgloss.AdaptiveColor{Light: "#2c2b2f", Dark: "#d3cdd4"}
subduedTextColor = lipgloss.AdaptiveColor{Light: "#454545", Dark: "#A49FA5"}
helpTextColor = lipgloss.AdaptiveColor{Light: "#DDDADA", Dark: "#3C3C3C"}
criticalThreatTextColor = lipgloss.AdaptiveColor{Light: "#FF0267", Dark: "#FF1E6F"} // "#ff1f7c" "#D2042D" "#eb2654"
highThreatColor = lipgloss.AdaptiveColor{Light: "#FF6C2D", Dark: "#FF6C2D"}
mediumThreatColor = lipgloss.AdaptiveColor{Light: "#B8860B", Dark: "#FFAF14"}
lowThreatColor = lipgloss.AdaptiveColor{Light: "#00A36C", Dark: "#00A36C"}
columnBorder = lipgloss.AdaptiveColor{Light: "#0BA4B8", Dark: "#AD58B4"} //"#c8c2d1"
separatorColor = lipgloss.AdaptiveColor{Light: "#0BA4B8", Dark: "#AD58B4"}
searchBorder = lipgloss.AdaptiveColor{Light: "#684eff", Dark: "#f792d4"}
separatorColor = lipgloss.AdaptiveColor{Light: "#0BA4B8", Dark: "#AD58B4"}
// catpuccin theme colors
red = lipgloss.AdaptiveColor{Light: "#D2042D", Dark: "#f38ba8"} // "#ff1f7c" "#D2042D" "#eb2654"
@@ -88,9 +79,9 @@ func (m *listModel) Init() tea.Cmd {
}
func (m *listModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg.(type) {
case tea.WindowSizeMsg:
// handle window resize
if _, ok := msg.(tea.WindowSizeMsg); ok {
_, v := listStyle.GetFrameSize()
m.Rows.SetSize(m.width, m.Rows.Height()-v)
}
@@ -99,7 +90,6 @@ func (m *listModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.Rows, cmd = m.Rows.Update(msg)
return m, cmd
}
func (m *listModel) SetHeight(height int) {
@@ -115,7 +105,7 @@ func (m *listModel) View() string {
header := renderColumnHeader(m.columns, m.width)
return listStyle.Copy().
return listStyle.
Border(lipgloss.RoundedBorder(), true, false, true, true).
BorderForeground(lavender).
Render(lipgloss.JoinVertical(lipgloss.Top, header, m.Rows.View()))
@@ -126,11 +116,10 @@ type listDelegate struct {
columns []column
}
func (d listDelegate) Height() int { return 2 }
func (d listDelegate) Spacing() int { return 1 }
func (d listDelegate) Update(msg tea.Msg, m *list.Model) tea.Cmd { return nil }
func (d listDelegate) Render(w io.Writer, m list.Model, index int, listItem list.Item) {
func (d listDelegate) Height() int { return 2 } //nolint:gocritic // bubbletea requires these to not be pointer methods
func (d listDelegate) Spacing() int { return 1 } //nolint:gocritic // bubbletea requires these to not be pointer methods
func (d listDelegate) Update(_ tea.Msg, _ *list.Model) tea.Cmd { return nil } //nolint:gocritic // bubbletea requires these to not be pointer methods
func (d listDelegate) Render(w io.Writer, m list.Model, index int, listItem list.Item) { //nolint:gocritic // bubbletea requires these to not be pointer methods
var (
severity string
@@ -141,24 +130,26 @@ func (d listDelegate) Render(w io.Writer, m list.Model, index int, listItem list
threatIntel string
)
if i, ok := listItem.(Item); ok {
severity = i.GetSeverity(true)
src = i.GetSrc()
dst = i.GetDst()
beacon = i.GetBeacon()
totalDuration = i.GetTotalDuration()
subdomains = i.GetSubdomains()
threatIntel = i.GetThreatIntel()
} else {
// get the item
i, ok := listItem.(*Item)
if !ok {
return
}
severity = i.GetSeverity(true)
src = i.GetSrc()
dst = i.GetDst()
beacon = i.GetBeacon()
totalDuration = i.GetTotalDuration()
subdomains = i.GetSubdomains()
threatIntel = i.GetThreatIntel()
if m.Width() <= 0 {
// short-circuit
return
}
// Conditions
// conditions
var (
isSelected = index == m.Index()
)
@@ -177,31 +168,31 @@ func (d listDelegate) Render(w io.Writer, m list.Model, index int, listItem list
}
// get severity
categoryStyle := style.Copy().PaddingLeft(2).Width(d.columns[0].width)
categoryTitle := categoryStyle.Render(Truncate(severity, categoryStyle))
categoryStyle := style.PaddingLeft(2).Width(d.columns[0].width)
categoryTitle := categoryStyle.Render(Truncate(severity, &categoryStyle))
// get source
srcStyle := style.Copy().Foreground(defaultTextColor).Width(d.columns[1].width)
srcTitle := srcStyle.Render(Truncate(src, srcStyle))
srcStyle := style.Foreground(defaultTextColor).Width(d.columns[1].width)
srcTitle := srcStyle.Render(Truncate(src, &srcStyle))
// get destination
dstStyle := style.Copy().Foreground(defaultTextColor).Width(d.columns[2].width)
dstTitle := dstStyle.Render(Truncate(dst, dstStyle))
dstStyle := style.Foreground(defaultTextColor).Width(d.columns[2].width)
dstTitle := dstStyle.Render(Truncate(dst, &dstStyle))
// get beacon
beaconStyle := style.Copy().Width(d.columns[3].width)
beaconStyle := style.Width(d.columns[3].width)
beaconTitle := beaconStyle.Render(beacon)
// get total duration
totalDurationStyle := style.Copy().Width(d.columns[4].width)
totalDurationStyle := style.Width(d.columns[4].width)
totalDurationTitle := totalDurationStyle.Render(totalDuration)
// get subdomains
subdomainsStyle := style.Copy().Width(d.columns[5].width)
subdomainsStyle := style.Width(d.columns[5].width)
subDomainsTitle := subdomainsStyle.Render(p.Sprint(subdomains))
// get threat intel
threatIntelStyle := style.Copy().Width(d.columns[6].width)
threatIntelStyle := style.Width(d.columns[6].width)
threatIntelTitle := threatIntelStyle.Render(p.Sprint(threatIntel))
// render the full row
@@ -215,7 +206,7 @@ func (d listDelegate) Render(w io.Writer, m list.Model, index int, listItem list
fmt.Fprintf(w, "%s", row)
}
func Truncate(str string, style lipgloss.Style) string {
func Truncate(str string, style *lipgloss.Style) string {
// Prevent text from exceeding list width
textwidth := uint(style.GetWidth() - style.GetPaddingLeft() - style.GetPaddingRight())
return truncate.StringWithTail(str, textwidth, ellipsis)
@@ -227,16 +218,16 @@ func renderIndicator(score float32, displayText string) string {
switch category {
case config.CriticalThreat:
return style.Copy().Foreground(red).Render(displayText)
return style.Foreground(red).Render(displayText)
case config.HighThreat:
return style.Copy().Foreground(peach).Render(displayText)
return style.Foreground(peach).Render(displayText)
case config.MediumThreat:
return style.Copy().Foreground(yellow).Render(displayText)
return style.Foreground(yellow).Render(displayText)
case config.LowThreat:
return style.Copy().Foreground(sapphire).Render(displayText)
return style.Foreground(sapphire).Render(displayText)
}
return style.Copy().Foreground(defaultTextColor).Render(displayText)
return style.Foreground(defaultTextColor).Render(displayText)
}
@@ -251,14 +242,14 @@ func renderColumnHeader(columns []column, headerWidth int) string {
// the fist column must start with a margin,
if i == 0 {
width -= 2 // subtract off the margin
header += columnStyle.Copy().MarginLeft(2).Width(width).Render(c.name)
header += columnStyle.MarginLeft(2).Width(width).Render(c.name)
} else {
header += columnStyle.Copy().Width(width).Render(c.name)
header += columnStyle.Width(width).Render(c.name)
}
// add a column border if not the last column
if i < len(columns)-1 {
header += columnStyle.Copy().Foreground(surface0).Render(" | ")
header += columnStyle.Foreground(surface0).Render(" | ")
}
}
+162 -9
View File
@@ -1,13 +1,166 @@
package viewer_test
// func testListStuff(t *testing.T) {
// var items []list.Item
// columns := []column{
// {"test", 10},
// }
import (
"github.com/activecm/rita/v5/viewer"
"github.com/stretchr/testify/require"
// listModel := MakeList(items, columns, 10, 10)
tea "github.com/charmbracelet/bubbletea"
)
// listModel.Init()
// listModel.Update(0)
// }
func (s *ViewerTestSuite) TestListScrolling() {
t := s.T()
// create new ui model
m, err := viewer.NewModel(s.maxTimestamp, s.minTimestamp, s.useCurrentTime, s.db)
require.NoError(t, err)
// get current selected index
initialSelectedIndex := m.List.Rows.Index()
// use down key to scroll the list down five times
for i := 0; i < 5; i++ {
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyDown,
},
))
}
// verify that the list was scrolled down five times from the initially selected index
require.Equal(t, initialSelectedIndex+5, m.List.Rows.Index())
// use up key to scroll the list up three times
for i := 0; i < 3; i++ {
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyUp,
},
))
}
// verify that the list was scrolled up 3 times, resulting in an index of 2 away from the initial index
require.Equal(t, initialSelectedIndex+2, m.List.Rows.Index())
}
func (s *ViewerTestSuite) TestListPaging() {
t := s.T()
// create new ui model
m, err := viewer.NewModel(s.maxTimestamp, s.minTimestamp, s.useCurrentTime, s.db)
require.NoError(t, err)
// get current page
initialPage := m.List.Rows.Paginator.Page
// select the 5th row in the list to ensure that the cursor is kept on the same row while paging
cursor := 5
m.List.Rows.Select(cursor)
// get current selected index
initialSelectedIndex := m.List.Rows.Index()
// use page down key to page down 5 pages in the list
for i := 0; i < 5; i++ {
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyPgDown,
},
))
}
// verify that the list was paged down five times from the initial page
require.Equal(t, initialPage+5, m.List.Rows.Paginator.Page, "after paging down 5 times, expected page to be %d, got %d", initialPage+5, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly
require.Equal(t, initialSelectedIndex+(m.List.Rows.Paginator.PerPage*5), m.List.Rows.Index(), "after paging down 5 times, expected selected index to be %d, got %d", initialSelectedIndex+(m.List.Rows.Paginator.PerPage*5), m.List.Rows.Index())
// verify that the cursor is still on the 5th row
require.Equal(t, cursor, m.List.Rows.Cursor(), "after paging down 5 times, expected cursor to remain as %d, got %d", cursor, m.List.Rows.Cursor())
// use page up key to scroll the list up 3 pages
for i := 0; i < 3; i++ {
m.Update(
tea.KeyMsg{
Type: tea.KeyPgUp,
},
)
}
// verify that the list was paged up 3 times, resulting in 2 pages away from the initial page (since we paged down 5 times first)
require.Equal(t, initialPage+2, m.List.Rows.Paginator.Page, "after paging up 3 times, expected page to be %d, got %d", initialPage+2, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly
require.Equal(t, initialSelectedIndex+(m.List.Rows.Paginator.PerPage*2), m.List.Rows.Index(), "after paging up 3 times, expected index to be %d, got %d", initialSelectedIndex+(m.List.Rows.Paginator.PerPage*2), m.List.Rows.Index())
// verify that the cursor is still on the correct row
require.Equal(t, cursor, m.List.Rows.Cursor(), "after paging up 3 times, expected cursor to remain as %d, got %d", cursor, m.List.Rows.Cursor())
// use home key to scroll to the start of the list
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyHome,
},
))
// verify that the list was paged to the start
require.Equal(t, 0, m.List.Rows.Paginator.Page, "after paging to the start, expected page to be %d, got %d", 0, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly
require.Equal(t, initialSelectedIndex, m.List.Rows.Index(), "after paging to the start, expected index to be %d, got %d", initialSelectedIndex, m.List.Rows.Index())
// verify that the cursor is still on the correct row
require.Equal(t, cursor, m.List.Rows.Cursor(), "after paging to the start, expected cursor to remain as %d, got %d", cursor, m.List.Rows.Cursor())
// use page end key to scroll to the end of the list
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyEnd,
},
))
// verify that the list was paged to the end
require.Equal(t, m.List.Rows.Paginator.TotalPages-1, m.List.Rows.Paginator.Page, "after paging to the end, expected page to be %d, got %d", m.List.Rows.Paginator.TotalPages-1, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly (since the last page may have fewer items than the cursor index, the selected index should be min(cursor, items on last page - 1 ))
endCursor := min(cursor, m.List.Rows.Paginator.ItemsOnPage(len(m.List.Rows.Items()))-1)
endIndex := (m.List.Rows.Paginator.Page * m.List.Rows.Paginator.PerPage) + endCursor
require.Equal(t, endIndex, m.List.Rows.Index(), "after paging to the end, expected selected index to be %d, got %d", endIndex, m.List.Rows.Index())
// verify that the cursor was updated to the correct spot
require.Equal(t, endCursor, m.List.Rows.Cursor(), "after paging to the end, expected cursor to be %d, got %d", endCursor, m.List.Rows.Cursor())
// page up one page (to the second-to-last page)
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyPgUp,
},
))
// set the cursor down to the bottom row of the second-to-last page to ensure that the selected row is greater than the number of items on the last page we will page to
m.List.Rows.Select((m.List.Rows.Paginator.Page * m.List.Rows.Paginator.PerPage) + (m.List.Rows.Paginator.ItemsOnPage(len(m.List.Rows.Items())) - 1))
// page down one page (back to last page)
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyPgDown,
},
))
// // verify that the list was paged up to the last page
require.Equal(t, m.List.Rows.Paginator.TotalPages-1, m.List.Rows.Paginator.Page, "after paging up to the last page, expected page to be %d, got %d", m.List.Rows.Paginator.TotalPages-1, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly to the last item on the last page
require.Equal(t, endIndex, m.List.Rows.Index(), "after paging up to the last page, expected selected index to be %d, got %d", endIndex, m.List.Rows.Index())
// verify that the cursor was updated to the correct spot
require.Equal(t, endCursor, m.List.Rows.Cursor(), "after paging up to the last page, expected cursor to be %d, got %d", endCursor, m.List.Rows.Cursor())
// fmt.Printf("page: %d\n", m.List.Rows.Paginator.Page)
// fmt.Printf("initialSelectedIndex: %d\n", initialSelectedIndex)
// fmt.Printf("actual index: %d\n", m.List.Rows.Index())
// fmt.Printf("perPage: %d\n", m.List.Rows.Paginator.PerPage)
// fmt.Printf("items on page: %d\n", m.List.Rows.Paginator.ItemsOnPage(len(m.List.Rows.Items())))
// fmt.Printf("cursor: %d\n", m.List.Rows.Cursor())
}
+16 -15
View File
@@ -50,29 +50,29 @@ type MixtapeResult struct {
type Item MixtapeResult
func (i Item) GetSrc() string {
func (i *Item) GetSrc() string {
if i.Src.String() == "::" && i.Dst.String() == "::" && len(i.FQDN) > 0 {
return ""
}
return i.Src.String()
}
func (i Item) GetDst() string {
func (i *Item) GetDst() string {
if i.Dst.String() == "::" && len(i.FQDN) > 0 {
return i.FQDN
}
return i.Dst.String()
}
// func (i item) FQDN() string { return i.fqdn }
func (i Item) GetBeacon() string {
// func (i *Item) FQDN() string { return i.fqdn }
func (i *Item) GetBeacon() string {
// if connection is a strobe, set beacon score to 100%
if i.StrobeScore > 0 {
return renderIndicator(i.StrobeScore, "100%")
}
return renderIndicator(i.BeaconThreatScore, fmt.Sprintf("%1.2f%%", i.BeaconScore*100))
}
func (i Item) GetFirstSeen(relativeTimestamp time.Time) string {
func (i *Item) GetFirstSeen(relativeTimestamp time.Time) string {
timeAgo := relativeTimestamp.Sub(i.FirstSeen)
switch {
case timeAgo.Hours() >= 8760:
@@ -111,27 +111,28 @@ func (i Item) GetFirstSeen(relativeTimestamp time.Time) string {
}
return fmt.Sprintf("%d %s ago", int(math.Floor(timeAgo.Hours())), text)
}
func (i Item) GetTotalDuration() string {
func (i *Item) GetTotalDuration() string {
return renderIndicator(i.LongConnScore, time.Duration(i.TotalDuration*float32(time.Second)).Truncate(time.Second).String())
}
func (i Item) GetPrevalence() string {
func (i *Item) GetPrevalence() string {
return renderIndicator(i.PrevalenceScore, fmt.Sprintf("%1.2f%%", i.Prevalence))
}
func (i Item) GetSubdomains() string {
func (i *Item) GetSubdomains() string {
return renderIndicator(i.C2OverDNSScore, fmt.Sprintf("%d", i.Subdomains))
}
func (i Item) GetPortProtoService() []string { return i.PortProtoService }
func (i *Item) GetPortProtoService() []string { return i.PortProtoService }
func (i Item) GetThreatIntel() string {
func (i *Item) GetThreatIntel() string {
if i.ThreatIntelScore > 0 {
return "⛔"
}
return ""
}
func (i Item) FilterValue() string { return i.GetSrc() } // no-op
func (i Item) GetSeverity(color bool) string {
// no-op
func (i Item) FilterValue() string { return i.GetSrc() } //nolint:gocritic // filtervalue cannot be a pointer method
func (i *Item) GetSeverity(color bool) string {
caser := cases.Title(language.English)
var severity config.ImpactCategory
@@ -156,7 +157,7 @@ func (i Item) GetSeverity(color bool) string {
return caser.String(string(severity))
}
func GetResults(db *database.DB, filter Filter, currentPage, pageSize int, minTimestamp time.Time) ([]list.Item, bool, error) {
func GetResults(db *database.DB, filter *Filter, currentPage, pageSize int, minTimestamp time.Time) ([]list.Item, bool, error) {
// build query
query, params, appliedFilter := BuildResultsQuery(filter, currentPage, pageSize, minTimestamp)
@@ -175,7 +176,7 @@ func GetResults(db *database.DB, filter Filter, currentPage, pageSize int, minTi
if err := rows.ScanStruct(&res); err != nil {
return nil, false, fmt.Errorf("could not read mixtape result for viewer: %w", err)
}
items = append(items, list.Item(res))
items = append(items, list.Item(&res))
}
rows.Close()
@@ -183,7 +184,7 @@ func GetResults(db *database.DB, filter Filter, currentPage, pageSize int, minTi
return items, appliedFilter, nil
}
func BuildResultsQuery(filter Filter, currentPage, pageSize int, minTimestamp time.Time) (string, clickhouse.Parameters, bool) {
func BuildResultsQuery(filter *Filter, currentPage, pageSize int, minTimestamp time.Time) (string, clickhouse.Parameters, bool) {
params := clickhouse.Parameters{}
query := `--sql
SELECT src, dst, fqdn,
+35 -51
View File
@@ -20,13 +20,6 @@ import (
var (
operatorRegex = regexp.MustCompile(`^(?P<operator>[><]=?)?(?P<value>(\d|[A-Za-z.])+)$`)
validSeverities = map[string]bool{
string(config.CriticalThreat): true,
string(config.HighThreat): true,
string(config.MediumThreat): true,
string(config.LowThreat): true,
}
allowedSortColumns = []string{"severity", "beacon", "duration", "subdomains"}
numericalColumns = []string{"count", "beacon", "subdomains"}
@@ -73,16 +66,11 @@ type searchModel struct {
}
func NewSearchModel(initialValue string, width int) searchModel {
// prompt := fmt.Sprintf(" is:%s ", sectionType)
// prompt := "hello world"
// prompt := "Search:"
ti := textinput.New()
ti.Placeholder = ""
ti.Focus()
// ti.Width = getInputWidth(width, prompt)
ti.PromptStyle = ti.PromptStyle.Copy().Foreground(mauve)
// ti.Prompt = prompt
ti.TextStyle = ti.TextStyle.Copy().Faint(true)
ti.PromptStyle = ti.PromptStyle.Foreground(mauve)
ti.TextStyle = ti.TextStyle.Faint(true)
ti.Blur()
ti.SetValue(initialValue)
ti.CursorStart()
@@ -91,21 +79,20 @@ func NewSearchModel(initialValue string, width int) searchModel {
TextInput: ti,
initialValue: initialValue,
width: width,
// sectionType: sectionType,
}
}
func (m searchModel) Init() tea.Cmd {
func (m *searchModel) Init() tea.Cmd {
return textinput.Blink
}
func (m searchModel) Update(msg tea.Msg) (searchModel, tea.Cmd) {
func (m *searchModel) Update(msg tea.Msg) (*searchModel, tea.Cmd) {
var cmd tea.Cmd
m.TextInput, cmd = m.TextInput.Update(msg)
return m, cmd
}
func (m searchModel) View() string {
func (m *searchModel) View() string {
helpStyle := lipgloss.NewStyle().Foreground(overlay0)
subduedHelpStyle := lipgloss.NewStyle().Foreground(surface0)
var label string
@@ -174,17 +161,17 @@ func (m searchModel) View() string {
}
func (m *searchModel) Focus() {
m.TextInput.TextStyle = m.TextInput.TextStyle.Copy().Faint(false)
m.TextInput.TextStyle = m.TextInput.TextStyle.Faint(false)
m.TextInput.CursorEnd()
m.TextInput.Focus()
}
func (m *searchModel) Blur() {
m.TextInput.TextStyle = m.TextInput.TextStyle.Copy().Faint(true)
m.TextInput.TextStyle = m.TextInput.TextStyle.Faint(true)
m.TextInput.Blur()
}
func (m searchModel) HasError() bool {
func (m *searchModel) HasError() bool {
return m.searchErr != ""
}
@@ -216,7 +203,7 @@ func (m *searchModel) ValidateSearchInput() {
}
}
func (m *searchModel) Filter() Filter {
func (m *searchModel) Filter() *Filter {
filter, err := ParseSearchInput(m.TextInput.Value())
if err != "" {
m.searchErr = err
@@ -225,18 +212,18 @@ func (m *searchModel) Filter() Filter {
}
// ParseSearchInput parses the search input and returns a filter struct
func ParseSearchInput(input string) (Filter, string) {
func ParseSearchInput(input string) (*Filter, string) {
// create a new filter struct
criteria := Filter{}
var criteria Filter
// return an empty filter if input is empty
if input == "" {
return Filter{}, ""
return nil, ""
}
// check for commas in the input
if strings.Contains(input, ",") {
return Filter{}, "commas are not supported"
return nil, "commas are not supported"
}
// split input into field-value pairs
@@ -250,7 +237,7 @@ func ParseSearchInput(input string) (Filter, string) {
// verify the search uses the proper colon-separated syntax
if !strings.Contains(input, ":") {
return Filter{}, "column name and value must be separated by a colon"
return nil, "column name and value must be separated by a colon"
}
// split the input into field and value
@@ -266,20 +253,20 @@ func ParseSearchInput(input string) (Filter, string) {
// parse operator and value from input
operator, number, parseErr := parseSearchOperator(field, value)
if parseErr != "" {
return Filter{}, parseErr
return nil, parseErr
}
// validate number is a true number
numberInt, err := strconv.Atoi(number)
if err != nil {
return Filter{}, field + " must be a valid number"
return nil, field + " must be a valid number"
}
// validate and format number to percentage (float from 0-1) for percentage columns
if slices.Contains(percentageColumns, field) {
// don't allow values over 100%
if numberInt > 100 {
return Filter{}, field + " can't be greater than 100"
return nil, field + " can't be greater than 100"
}
percentage := float32(numberInt) / 100
// divide value by 100 to convert to decimal, must be formatted to two places
@@ -314,7 +301,7 @@ func ParseSearchInput(input string) (Filter, string) {
// parse operator and time string from value
operator, input, parseErr := parseSearchOperator(field, value)
if parseErr != "" {
return Filter{}, parseErr
return nil, parseErr
}
fmt.Println("operator:", operator, "input", input, "parseErr", parseErr)
@@ -323,7 +310,7 @@ func ParseSearchInput(input string) (Filter, string) {
duration, err := time.ParseDuration(input)
if err != nil {
parseErr = field + " must be a valid time in the format '10s', '1.5h', '2h45m', etc. Valid units are 's', 'm', 'h'"
return Filter{}, parseErr
return nil, parseErr
}
// assign operator to criteria
@@ -343,7 +330,7 @@ func ParseSearchInput(input string) (Filter, string) {
case "src":
// validate string is IP address
if _, err := netip.ParseAddr(value); err != nil {
return Filter{}, "src must be a valid IP address"
return nil, "src must be a valid IP address"
}
criteria.Src = value
@@ -351,18 +338,19 @@ func ParseSearchInput(input string) (Filter, string) {
// validate if string is IP address
if _, err := netip.ParseAddr(value); err != nil {
// if value is not an IP, check for valid FQDN
if util.ValidFQDN(value) {
criteria.Fqdn = value
} else {
return Filter{}, "dst must be a valid IP address or FQDN"
if !util.ValidFQDN(value) {
return nil, "dst must be a valid IP address or FQDN"
}
criteria.Fqdn = value
} else {
criteria.Dst = value
}
case "threat_intel":
filter, err := strconv.ParseBool(value)
if err != nil {
return Filter{}, "threat_intel must be true or false"
return nil, "threat_intel must be true or false"
}
if filter {
criteria.ThreatIntel = "true"
@@ -373,7 +361,7 @@ func ParseSearchInput(input string) (Filter, string) {
// split the column from the sort direction
sortSplit := strings.Split(value, "-")
if len(sortSplit) != 2 {
return Filter{}, "sort value must contain one hyphen, in the format sort:<column>-<direction>"
return nil, "sort value must contain one hyphen, in the format sort:<column>-<direction>"
}
// validate sort column and direction
@@ -382,12 +370,12 @@ func ParseSearchInput(input string) (Filter, string) {
// make sure this column has sorting enabled
if !slices.Contains(allowedSortColumns, column) {
return Filter{}, "invalid sort column"
return nil, "invalid sort column"
}
// validate sort direction
if direction != "asc" && direction != "desc" {
return Filter{}, "sort direction must be either asc or desc"
return nil, "sort direction must be either asc or desc"
}
// assign sort column and direction to criteria
@@ -413,7 +401,7 @@ func ParseSearchInput(input string) (Filter, string) {
// via modifiers, so users aren't able to set assign threat categories in the config to critical
// We also do not want to allow users to filter by the None category
if err != nil && category != config.CriticalThreat || category == config.NoneThreat {
return Filter{}, "invalid category, must be 'critical', 'high', 'medium', or 'low'"
return nil, "invalid category, must be 'critical', 'high', 'medium', or 'low'"
}
// assign severity to criteria and set the needed operators for querying
@@ -427,8 +415,7 @@ func ParseSearchInput(input string) (Filter, string) {
criteria.Severity = append(criteria.Severity, OperatorFilter{
Operator: "<=",
Value: fmt.Sprint(config.HIGH_CATEGORY_SCORE),
})
criteria.Severity = append(criteria.Severity, OperatorFilter{
}, OperatorFilter{
Operator: ">=",
Value: fmt.Sprint(config.MEDIUM_CATEGORY_SCORE),
})
@@ -436,8 +423,7 @@ func ParseSearchInput(input string) (Filter, string) {
criteria.Severity = append(criteria.Severity, OperatorFilter{
Operator: "<",
Value: fmt.Sprint(config.MEDIUM_CATEGORY_SCORE),
})
criteria.Severity = append(criteria.Severity, OperatorFilter{
}, OperatorFilter{
Operator: ">=",
Value: fmt.Sprint(config.LOW_CATEGORY_SCORE),
})
@@ -445,8 +431,7 @@ func ParseSearchInput(input string) (Filter, string) {
criteria.Severity = append(criteria.Severity, OperatorFilter{
Operator: "<",
Value: fmt.Sprint(config.LOW_CATEGORY_SCORE),
})
criteria.Severity = append(criteria.Severity, OperatorFilter{
}, OperatorFilter{
Operator: ">=",
Value: fmt.Sprint(config.NONE_CATEGORY_SCORE),
})
@@ -454,13 +439,12 @@ func ParseSearchInput(input string) (Filter, string) {
}
default:
return Filter{}, "please reference a valid search column"
return nil, "please reference a valid search column"
}
}
// panic(fmt.Sprintf("end of function criteria: %+v", criteria))
return criteria, ""
return &criteria, ""
}
func parseSearchOperator(field string, value string) (string, string, string) {
+41 -41
View File
@@ -171,72 +171,72 @@ func (s *ViewerTestSuite) TestSearchResults() {
type testCase struct {
name string
filter viewer.Filter
valid func(viewer.Item) bool
sorted func(float64, viewer.Item) (float64, bool) // return whether or not the next item follows the right sort order
field func(viewer.Item) float64 // returns the field of the column being sorted
valid func(*viewer.Item) bool
sorted func(float64, *viewer.Item) (float64, bool) // return whether or not the next item follows the right sort order
field func(*viewer.Item) float64 // returns the field of the column being sorted
checkSorting bool
}
cases := []testCase{
{name: "Filter by src IP", filter: viewer.Filter{Src: "10.55.100.100"}, valid: func(i viewer.Item) bool { return i.Src.String() == "10.55.100.100" }},
{name: "Filter by dst IP", filter: viewer.Filter{Dst: "165.227.88.15"}, valid: func(i viewer.Item) bool { return i.Dst.String() == "165.227.88.15" }},
{name: "Filter by FQDN", filter: viewer.Filter{Fqdn: "www.alexa.com"}, valid: func(i viewer.Item) bool { return i.FQDN == "www.alexa.com" }},
{name: "Filter by src IP", filter: viewer.Filter{Src: "10.55.100.100"}, valid: func(i *viewer.Item) bool { return i.Src.String() == "10.55.100.100" }},
{name: "Filter by dst IP", filter: viewer.Filter{Dst: "165.227.88.15"}, valid: func(i *viewer.Item) bool { return i.Dst.String() == "165.227.88.15" }},
{name: "Filter by FQDN", filter: viewer.Filter{Fqdn: "www.alexa.com"}, valid: func(i *viewer.Item) bool { return i.FQDN == "www.alexa.com" }},
// beacon
{name: "Filter by beacon score", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: "=", Value: "1"}}, valid: func(i viewer.Item) bool { return i.BeaconScore == 1 }},
{name: "Filter by beacon score, greater than", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: ">", Value: "0.98"}, SortBeacon: "asc"}, valid: func(i viewer.Item) bool { return i.BeaconScore > 0.98 }},
{name: "Filter by beacon score, greater than or equal", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: ">=", Value: "0.98"}}, valid: func(i viewer.Item) bool { return i.BeaconScore >= 0.98 }},
{name: "Filter by beacon score, greater than or equal", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: ">=", Value: "0.98"}}, valid: func(i viewer.Item) bool { return i.BeaconScore >= 0.98 }},
{name: "Filter by beacon score, less than", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: "<", Value: "0.70"}}, valid: func(i viewer.Item) bool { return i.BeaconScore < 0.7 }},
{name: "Filter by beacon score, less than or equal", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: "<=", Value: "0.70"}}, valid: func(i viewer.Item) bool { return i.BeaconScore <= 0.7 }},
{name: "Filter by beacon score", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: "=", Value: "1"}}, valid: func(i *viewer.Item) bool { return i.BeaconScore == 1 }},
{name: "Filter by beacon score, greater than", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: ">", Value: "0.98"}, SortBeacon: "asc"}, valid: func(i *viewer.Item) bool { return i.BeaconScore > 0.98 }},
{name: "Filter by beacon score, greater than or equal", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: ">=", Value: "0.98"}}, valid: func(i *viewer.Item) bool { return i.BeaconScore >= 0.98 }},
{name: "Filter by beacon score, greater than or equal", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: ">=", Value: "0.98"}}, valid: func(i *viewer.Item) bool { return i.BeaconScore >= 0.98 }},
{name: "Filter by beacon score, less than", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: "<", Value: "0.70"}}, valid: func(i *viewer.Item) bool { return i.BeaconScore < 0.7 }},
{name: "Filter by beacon score, less than or equal", filter: viewer.Filter{Beacon: viewer.OperatorFilter{Operator: "<=", Value: "0.70"}}, valid: func(i *viewer.Item) bool { return i.BeaconScore <= 0.7 }},
// duration
{name: "Filter by duration", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: "=", Value: "2584"}}, valid: func(i viewer.Item) bool { return math.Floor(float64(i.TotalDuration)) == 2584 }},
{name: "Filter by duration, greater than", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: ">", Value: "21600"}}, valid: func(i viewer.Item) bool { return i.TotalDuration > 21600 }},
{name: "Filter by duration, greater than or equal", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: ">=", Value: "21600"}}, valid: func(i viewer.Item) bool { return i.TotalDuration >= 21600 }},
{name: "Filter by duration, less than", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: "<", Value: "3600"}}, valid: func(i viewer.Item) bool { return i.TotalDuration < 3600 }},
{name: "Filter by duration, less than or equal", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: "<=", Value: "3600"}}, valid: func(i viewer.Item) bool { return i.TotalDuration <= 3600 }},
{name: "Filter by duration", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: "=", Value: "2584"}}, valid: func(i *viewer.Item) bool { return math.Floor(float64(i.TotalDuration)) == 2584 }},
{name: "Filter by duration, greater than", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: ">", Value: "21600"}}, valid: func(i *viewer.Item) bool { return i.TotalDuration > 21600 }},
{name: "Filter by duration, greater than or equal", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: ">=", Value: "21600"}}, valid: func(i *viewer.Item) bool { return i.TotalDuration >= 21600 }},
{name: "Filter by duration, less than", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: "<", Value: "3600"}}, valid: func(i *viewer.Item) bool { return i.TotalDuration < 3600 }},
{name: "Filter by duration, less than or equal", filter: viewer.Filter{Duration: viewer.OperatorFilter{Operator: "<=", Value: "3600"}}, valid: func(i *viewer.Item) bool { return i.TotalDuration <= 3600 }},
// subdomains
{name: "Filter by subdomains", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: "=", Value: "62468"}}, valid: func(i viewer.Item) bool { return i.Subdomains == 62468 }},
{name: "Filter by subdomains, greater than", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: ">", Value: "50"}}, valid: func(i viewer.Item) bool { return i.Subdomains > 50 }},
{name: "Filter by subdomains, greater than or equal", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: ">=", Value: "50"}}, valid: func(i viewer.Item) bool { return i.Subdomains >= 50 }},
{name: "Filter by subdomains, less than", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: "<", Value: "100"}}, valid: func(i viewer.Item) bool { return i.Subdomains < 100 }},
{name: "Filter by subdomains, less than or equal", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: "<=", Value: "100"}}, valid: func(i viewer.Item) bool { return i.Subdomains <= 100 }},
{name: "Filter by subdomains", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: "=", Value: "62468"}}, valid: func(i *viewer.Item) bool { return i.Subdomains == 62468 }},
{name: "Filter by subdomains, greater than", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: ">", Value: "50"}}, valid: func(i *viewer.Item) bool { return i.Subdomains > 50 }},
{name: "Filter by subdomains, greater than or equal", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: ">=", Value: "50"}}, valid: func(i *viewer.Item) bool { return i.Subdomains >= 50 }},
{name: "Filter by subdomains, less than", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: "<", Value: "100"}}, valid: func(i *viewer.Item) bool { return i.Subdomains < 100 }},
{name: "Filter by subdomains, less than or equal", filter: viewer.Filter{Subdomains: viewer.OperatorFilter{Operator: "<=", Value: "100"}}, valid: func(i *viewer.Item) bool { return i.Subdomains <= 100 }},
// threat intel
// {name: "Filter by threat intel, true", filter: viewer.Filter{ThreatIntel: "true"}, valid: func(i viewer.Item) bool { return i.ThreatIntelScore > 0 }},
{name: "Filter by threat intel, false", filter: viewer.Filter{ThreatIntel: "false"}, valid: func(i viewer.Item) bool { return i.ThreatIntelScore == 0 }},
// {name: "Filter by threat intel, true", filter: viewer.Filter{ThreatIntel: "true"}, valid: func(i *viewer.Item) bool { return i.ThreatIntelScore > 0 }},
{name: "Filter by threat intel, false", filter: viewer.Filter{ThreatIntel: "false"}, valid: func(i *viewer.Item) bool { return i.ThreatIntelScore == 0 }},
// severity
{name: "Filter by severity, critical", filter: viewer.Filter{Severity: []viewer.OperatorFilter{{Operator: ">", Value: fmt.Sprint(config.HIGH_CATEGORY_SCORE)}}}, valid: func(i viewer.Item) bool { return i.FinalScore > config.HIGH_CATEGORY_SCORE }},
{name: "Filter by severity, high", filter: viewer.Filter{Severity: []viewer.OperatorFilter{{Operator: "<=", Value: fmt.Sprint(config.HIGH_CATEGORY_SCORE)}, {Operator: ">=", Value: fmt.Sprint(config.MEDIUM_CATEGORY_SCORE)}}}, valid: func(i viewer.Item) bool {
{name: "Filter by severity, critical", filter: viewer.Filter{Severity: []viewer.OperatorFilter{{Operator: ">", Value: fmt.Sprint(config.HIGH_CATEGORY_SCORE)}}}, valid: func(i *viewer.Item) bool { return i.FinalScore > config.HIGH_CATEGORY_SCORE }},
{name: "Filter by severity, high", filter: viewer.Filter{Severity: []viewer.OperatorFilter{{Operator: "<=", Value: fmt.Sprint(config.HIGH_CATEGORY_SCORE)}, {Operator: ">=", Value: fmt.Sprint(config.MEDIUM_CATEGORY_SCORE)}}}, valid: func(i *viewer.Item) bool {
return i.FinalScore <= config.HIGH_CATEGORY_SCORE && i.FinalScore >= config.MEDIUM_CATEGORY_SCORE
}},
{name: "Filter by severity, medium", filter: viewer.Filter{Severity: []viewer.OperatorFilter{{Operator: "<", Value: fmt.Sprint(config.MEDIUM_CATEGORY_SCORE)}, {Operator: ">=", Value: fmt.Sprint(config.LOW_CATEGORY_SCORE)}}}, valid: func(i viewer.Item) bool {
{name: "Filter by severity, medium", filter: viewer.Filter{Severity: []viewer.OperatorFilter{{Operator: "<", Value: fmt.Sprint(config.MEDIUM_CATEGORY_SCORE)}, {Operator: ">=", Value: fmt.Sprint(config.LOW_CATEGORY_SCORE)}}}, valid: func(i *viewer.Item) bool {
return i.FinalScore < config.MEDIUM_CATEGORY_SCORE && i.FinalScore >= config.LOW_CATEGORY_SCORE
}},
{name: "Filter by severity, low", filter: viewer.Filter{Severity: []viewer.OperatorFilter{{Operator: "<", Value: fmt.Sprint(config.LOW_CATEGORY_SCORE)}, {Operator: ">=", Value: fmt.Sprint(config.NONE_CATEGORY_SCORE)}}}, valid: func(i viewer.Item) bool {
{name: "Filter by severity, low", filter: viewer.Filter{Severity: []viewer.OperatorFilter{{Operator: "<", Value: fmt.Sprint(config.LOW_CATEGORY_SCORE)}, {Operator: ">=", Value: fmt.Sprint(config.NONE_CATEGORY_SCORE)}}}, valid: func(i *viewer.Item) bool {
return i.FinalScore < config.LOW_CATEGORY_SCORE && i.FinalScore >= config.NONE_CATEGORY_SCORE
}},
// sorting
{name: "Sort by beacon, desc", filter: viewer.Filter{SortBeacon: "desc"}, checkSorting: true, field: func(item viewer.Item) float64 { return float64(item.BeaconScore) }, sorted: func(currentVal float64, newItem viewer.Item) (float64, bool) {
{name: "Sort by beacon, desc", filter: viewer.Filter{SortBeacon: "desc"}, checkSorting: true, field: func(item *viewer.Item) float64 { return float64(item.BeaconScore) }, sorted: func(currentVal float64, newItem *viewer.Item) (float64, bool) {
return float64(newItem.BeaconScore), newItem.BeaconScore <= float32(currentVal)
}},
{name: "Sort by beacon, asc", filter: viewer.Filter{SortBeacon: "asc"}, checkSorting: true, field: func(item viewer.Item) float64 { return float64(item.BeaconScore) }, sorted: func(currentVal float64, newItem viewer.Item) (float64, bool) {
{name: "Sort by beacon, asc", filter: viewer.Filter{SortBeacon: "asc"}, checkSorting: true, field: func(item *viewer.Item) float64 { return float64(item.BeaconScore) }, sorted: func(currentVal float64, newItem *viewer.Item) (float64, bool) {
return float64(newItem.BeaconScore), newItem.BeaconScore >= float32(currentVal)
}},
{name: "Sort by duration, desc", filter: viewer.Filter{SortDuration: "desc"}, checkSorting: true, field: func(item viewer.Item) float64 { return float64(item.TotalDuration) }, sorted: func(currentVal float64, newItem viewer.Item) (float64, bool) {
{name: "Sort by duration, desc", filter: viewer.Filter{SortDuration: "desc"}, checkSorting: true, field: func(item *viewer.Item) float64 { return float64(item.TotalDuration) }, sorted: func(currentVal float64, newItem *viewer.Item) (float64, bool) {
return float64(newItem.TotalDuration), newItem.TotalDuration <= float32(currentVal)
}},
{name: "Sort by duration, asc", filter: viewer.Filter{SortDuration: "asc"}, checkSorting: true, field: func(item viewer.Item) float64 { return float64(item.TotalDuration) }, sorted: func(currentVal float64, newItem viewer.Item) (float64, bool) {
{name: "Sort by duration, asc", filter: viewer.Filter{SortDuration: "asc"}, checkSorting: true, field: func(item *viewer.Item) float64 { return float64(item.TotalDuration) }, sorted: func(currentVal float64, newItem *viewer.Item) (float64, bool) {
return float64(newItem.TotalDuration), newItem.TotalDuration >= float32(currentVal)
}},
{name: "Sort by severity, desc", filter: viewer.Filter{SortSeverity: "desc"}, checkSorting: true, field: func(item viewer.Item) float64 { return float64(item.FinalScore) }, sorted: func(currentVal float64, newItem viewer.Item) (float64, bool) {
{name: "Sort by severity, desc", filter: viewer.Filter{SortSeverity: "desc"}, checkSorting: true, field: func(item *viewer.Item) float64 { return float64(item.FinalScore) }, sorted: func(currentVal float64, newItem *viewer.Item) (float64, bool) {
return float64(newItem.FinalScore), newItem.FinalScore <= float32(currentVal)
}},
{name: "Sort by severity, asc", filter: viewer.Filter{SortSeverity: "asc"}, checkSorting: true, field: func(item viewer.Item) float64 { return float64(item.FinalScore) }, sorted: func(currentVal float64, newItem viewer.Item) (float64, bool) {
{name: "Sort by severity, asc", filter: viewer.Filter{SortSeverity: "asc"}, checkSorting: true, field: func(item *viewer.Item) float64 { return float64(item.FinalScore) }, sorted: func(currentVal float64, newItem *viewer.Item) (float64, bool) {
return float64(newItem.FinalScore), newItem.FinalScore >= float32(currentVal)
}},
{name: "Sort by subdomains, desc", filter: viewer.Filter{SortSubdomains: "desc"}, checkSorting: true, field: func(item viewer.Item) float64 { return float64(item.Subdomains) }, sorted: func(currentVal float64, newItem viewer.Item) (float64, bool) {
{name: "Sort by subdomains, desc", filter: viewer.Filter{SortSubdomains: "desc"}, checkSorting: true, field: func(item *viewer.Item) float64 { return float64(item.Subdomains) }, sorted: func(currentVal float64, newItem *viewer.Item) (float64, bool) {
return float64(newItem.Subdomains), float64(newItem.Subdomains) <= currentVal
}},
{name: "Sort by subdomains, asc", filter: viewer.Filter{SortSubdomains: "asc"}, checkSorting: true, field: func(item viewer.Item) float64 { return float64(item.Subdomains) }, sorted: func(currentVal float64, newItem viewer.Item) (float64, bool) {
{name: "Sort by subdomains, asc", filter: viewer.Filter{SortSubdomains: "asc"}, checkSorting: true, field: func(item *viewer.Item) float64 { return float64(item.Subdomains) }, sorted: func(currentVal float64, newItem *viewer.Item) (float64, bool) {
return float64(newItem.Subdomains), float64(newItem.Subdomains) >= currentVal
}},
}
@@ -244,7 +244,7 @@ func (s *ViewerTestSuite) TestSearchResults() {
test := cases[i]
s.Run(test.name, func() {
// get filter from search bar
res, appliedFilter, err := viewer.GetResults(s.db, test.filter, 0, 20, s.minTimestamp)
res, appliedFilter, err := viewer.GetResults(s.db, &test.filter, 0, 20, s.minTimestamp)
require.NoError(t, err)
require.True(t, appliedFilter, "filter criteria must be applied")
require.NotEmpty(t, res, "results should not be empty")
@@ -256,7 +256,7 @@ func (s *ViewerTestSuite) TestSearchResults() {
// check that the results match the search criteria
valid := true
for _, r := range res {
valid = test.valid(r.(viewer.Item))
valid = test.valid(r.(*viewer.Item))
}
require.True(t, valid, "all results should match the search criteria")
}
@@ -265,14 +265,14 @@ func (s *ViewerTestSuite) TestSearchResults() {
}
// validateSorting checks whether or not results are sorted by a particular column
func validateSorting(items []list.Item, field func(viewer.Item) float64, sorted func(float64, viewer.Item) (float64, bool)) bool {
func validateSorting(items []list.Item, field func(*viewer.Item) float64, sorted func(float64, *viewer.Item) (float64, bool)) bool {
var current float64
for i, item := range items {
if i == 0 {
// set the initial value by getting the right field for the first item
current = field(item.(viewer.Item))
current = field(item.(*viewer.Item))
}
res, ok := item.(viewer.Item)
res, ok := item.(*viewer.Item)
if !ok {
return false
}
+6 -82
View File
@@ -22,13 +22,13 @@ type modifier struct {
type sidebarModel struct {
Viewport viewport.Model
Data Item
Data *Item
maxTimestamp time.Time
useCurrentTime bool
ScrollEnabled bool
}
func NewSidebarModel(maxTS time.Time, useCurrentTime bool, initialData Item) sidebarModel {
func NewSidebarModel(maxTS time.Time, useCurrentTime bool, initialData *Item) sidebarModel {
return sidebarModel{
Viewport: viewport.Model{},
maxTimestamp: maxTS,
@@ -41,18 +41,7 @@ func (m *sidebarModel) Init() tea.Cmd {
m.Viewport.SetContent(m.getSidebarContents())
return nil
}
func (m *sidebarModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
// switch msg := msg.(type) {
// case tea.KeyMsg:
// switch {
// // case key.Matches(msg, keys.Keys.PageDown):
// // m.viewport.HalfViewDown()
// // case key.Matches(msg, keys.Keys.PageUp):
// // m.viewport.HalfViewUp()
// }
// }
func (m *sidebarModel) Update(_ tea.Msg) (tea.Model, tea.Cmd) {
return m, nil
}
@@ -86,7 +75,7 @@ func (m *sidebarModel) getSidebarContents() string {
fqdnLabel := "FQDN"
dstStyle := lipgloss.NewStyle().Width(m.Viewport.Width - len(fqdnLabel) - (headerPadding * 4))
target = lipgloss.JoinHorizontal(lipgloss.Left, headerLabelStyle.Render(fqdnLabel), headerValueStyle.Copy().Render(Truncate(m.Data.GetDst(), dstStyle)))
target = lipgloss.JoinHorizontal(lipgloss.Left, headerLabelStyle.Render(fqdnLabel), headerValueStyle.Render(Truncate(m.Data.GetDst(), &dstStyle)))
target = lipgloss.NewStyle().MarginBottom(2).Render(target)
} else {
// handle connection pair, ip -> ip or ip -> fqdn
@@ -96,8 +85,8 @@ func (m *sidebarModel) getSidebarContents() string {
dstStyle := lipgloss.NewStyle().Width(m.Viewport.Width - len(dstLabel) - (headerPadding * 4))
// - 6 - len(m.data.GetSrc())
// dstStyle := lipgloss.NewStyle().Width(dstWidth)
src := lipgloss.JoinHorizontal(lipgloss.Left, headerLabelStyle.Render(srcLabel), headerValueStyle.Copy().Render(Truncate(m.Data.GetSrc(), srcStyle)))
dst := lipgloss.JoinHorizontal(lipgloss.Left, headerLabelStyle.Render(dstLabel), headerValueStyle.Copy().Render(Truncate(m.Data.GetDst(), dstStyle)))
src := lipgloss.JoinHorizontal(lipgloss.Left, headerLabelStyle.Render(srcLabel), headerValueStyle.Render(Truncate(m.Data.GetSrc(), &srcStyle)))
dst := lipgloss.JoinHorizontal(lipgloss.Left, headerLabelStyle.Render(dstLabel), headerValueStyle.Render(Truncate(m.Data.GetDst(), &dstStyle)))
target = lipgloss.JoinVertical(lipgloss.Top, lipgloss.NewStyle().MarginBottom(1).Render(src), dst)
}
heading := lipgloss.NewStyle().
@@ -271,70 +260,5 @@ func renderModifier(mod modifier) string {
data := lipgloss.NewStyle().Foreground(defaultTextColor).Render(mod.value)
modifier := lipgloss.JoinVertical(lipgloss.Top, header, data)
// modifier = lipgloss.NewStyle().Border(lipgloss.NormalBorder(), false, true, false, false).Padding(0, 2, 0, 0).Render(modifier)
return modifier
}
func renderPorts(portProtoService []string, viewportWidth int, remainingLines int) string {
if len(portProtoService) == 0 {
return ""
}
// create style for header
portsHeaderStyle := lipgloss.NewStyle().Background(overlay2).Foreground(base).Bold(true).Padding(0, 2).MarginTop(1)
// render header
portsHeader := portsHeaderStyle.Render("Port : Proto : Service")
// calculate number of items each column should have
itemsPerColumn := remainingLines - lipgloss.Height(portsHeader)
if itemsPerColumn < 1 {
itemsPerColumn = 1
}
// calculate the number of columns needed
numColumns := (len(portProtoService) + itemsPerColumn - 1) / itemsPerColumn
// determine the maximum width per column
estimatedColumnWidth := lipgloss.Width(portProtoService[0]) + 3
// calculate the number of columns that can be displayed on the screen
numDisplayColumns := viewportWidth / estimatedColumnWidth
// create style for data
dataStyle := lipgloss.NewStyle().Width(estimatedColumnWidth)
// create columns
columns := make([][]string, numColumns)
for i := range columns {
columnStart := i * itemsPerColumn
columnEnd := columnStart + itemsPerColumn
if columnEnd > len(portProtoService) {
columnEnd = len(portProtoService)
}
columns[i] = portProtoService[columnStart:columnEnd]
}
// render columns
columnStrs := make([]string, 0)
for i := 0; i < numDisplayColumns && i < numColumns; i++ {
// get column
col := columns[i]
// replace last item in column with ellipsis if there are more columns than can be displayed
if i == numDisplayColumns-1 && numDisplayColumns < numColumns {
col[len(col)-1] = "..."
}
// add column to list of columns to be rendered
columnStrs = append(columnStrs, dataStyle.Render(lipgloss.JoinVertical(lipgloss.Left, col...)))
}
// join columns horizontally
portsData := lipgloss.JoinHorizontal(lipgloss.Top, columnStrs...)
// combine the header with data
ports := lipgloss.JoinVertical(lipgloss.Top, portsHeader, portsData)
return ports
}
+31 -24
View File
@@ -3,7 +3,6 @@ package viewer
import (
"fmt"
"math"
"os"
"runtime"
"slices"
"time"
@@ -25,7 +24,7 @@ var mainStyle = lipgloss.NewStyle().Margin(0, 0)
type Model struct {
// keys keys.KeyMap
minTS time.Time
SearchBar searchModel
SearchBar *searchModel
SideBar sidebarModel
List listModel
searchValue string
@@ -59,17 +58,19 @@ type column struct {
}
// CreateUI creates the terminal UI
func CreateUI(cfg *config.Config, db *database.DB, useCurrentTime bool, maxTimestamp time.Time, minTimestamp time.Time) error {
func CreateUI(_ *config.Config, db *database.DB, useCurrentTime bool, maxTimestamp time.Time, minTimestamp time.Time) error {
// create model
m, err := NewModel(maxTimestamp, minTimestamp, useCurrentTime, db)
if err != nil {
return err
}
// create program
p := tea.NewProgram(m, tea.WithAltScreen())
// run the program
if _, err := p.Run(); err != nil {
fmt.Println("Error running program:", err)
os.Exit(1)
return fmt.Errorf("error running program: %w", err)
}
return nil
@@ -78,38 +79,44 @@ func CreateUI(cfg *config.Config, db *database.DB, useCurrentTime bool, maxTimes
func NewModel(maxTimestamp, minTimestamp time.Time, useCurrentTime bool, db *database.DB) (*Model, error) {
pageSize := 100
// get results from database
rows, _, err := GetResults(db, Filter{}, 0, pageSize, minTimestamp)
rows, _, err := GetResults(db, &Filter{}, 0, pageSize, minTimestamp)
if err != nil {
return nil, err
}
// set columns
columns := []column{{"Severity", 14}, {"Source", 20}, {"Destination", 30}, {"Beacon", 10}, {"Duration", 15}, {"Subdomains", 15}, {"Threat Intel", 15}}
// set table size
width := getTableWidth(columns)
height := 20
// create list
list := MakeList(rows, columns, width, height)
// create dataList
dataList := MakeList(rows, columns, width, height)
// create search bar
searchBar := NewSearchModel("", width)
// create side bar
sideBar := NewSidebarModel(maxTimestamp, useCurrentTime, Item{})
if len(list.Rows.Items()) > 0 {
if data, ok := list.Rows.Items()[list.Rows.Index()].(Item); ok {
sideBar.Data = data
} else {
sideBar := NewSidebarModel(maxTimestamp, useCurrentTime, &Item{})
if len(dataList.Rows.Items()) > 0 {
// set sidebar data to whichever item is selected in the list
data, ok := dataList.Rows.Items()[dataList.Rows.Index()].(Item)
if !ok {
return nil, fmt.Errorf("error setting sidebar data")
}
sideBar.Data = &data
}
// create footer
footer := NewFooterModel(db.GetSelectedDB())
// create model
m := &Model{
minTS: minTimestamp,
List: list,
SearchBar: searchBar,
List: dataList,
SearchBar: &searchBar,
SideBar: sideBar,
serverPageSize: pageSize,
Footer: footer,
@@ -117,8 +124,10 @@ func NewModel(maxTimestamp, minTimestamp time.Time, useCurrentTime bool, db *dat
width: width,
}
// initialize model components
m.Init()
// initialize sidebar
m.SideBar.Init()
// create terminal ui model
@@ -126,8 +135,6 @@ func NewModel(maxTimestamp, minTimestamp time.Time, useCurrentTime bool, db *dat
}
func (m *Model) Init() tea.Cmd {
// set footer to search help text
// m.footer = searchHelpText()
// set title
m.title = getTitle()
@@ -178,7 +185,7 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
case tea.WindowSizeMsg:
// make the footer the entire width of the terminal
m.Footer.width = msg.Width
list := m.List.View()
// make the list fill the extra vertical space
m.List.SetHeight(msg.Height - int(math.Max(float64(lipgloss.Height(m.SearchBar.View())), float64(lipgloss.Height(m.title)))) - lipgloss.Height(m.dbFooterBar))
@@ -186,7 +193,7 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.SideBar.Viewport.Height = m.List.totalHeight
// make sidebar fill the extra horizontal space
m.SideBar.Viewport.Width = msg.Width - lipgloss.Width(list) - 4
m.SideBar.Viewport.Width = msg.Width - lipgloss.Width(m.List.View()) - 4
// make search bar the same width as the list
m.SearchBar.width = m.List.width
@@ -252,11 +259,11 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
}
if data, ok := m.List.Rows.Items()[m.List.Rows.Index()].(Item); ok {
m.SideBar.Data = data
m.SideBar.Data = &data
}
} else {
m.SideBar.Data = Item{}
m.SideBar.Data = &Item{}
}
return m, cmd
@@ -275,7 +282,7 @@ func (m *Model) View() string {
default:
mainContent = lipgloss.JoinHorizontal(
lipgloss.Left,
mainStyle.Copy().Render(m.List.View()),
mainStyle.Render(m.List.View()),
mainStyle.Render(m.SideBar.View()),
)
}
@@ -445,7 +452,7 @@ func (m *Model) resetFiltering() {
// getTitle returns the title of the application
func getTitle() string {
return mainStyle.Copy().
return mainStyle.
MarginLeft(1).MarginTop(3).
// DO NOT INDENT THE FOLLOWING CODE BLOCK!
Render(`
@@ -588,7 +595,7 @@ func mainHelpText() string {
}
func helpPanel(height int, width int, contents string) string {
return mainStyle.Copy().Height(height).Width(width).
return mainStyle.Height(height).Width(width).
Border(lipgloss.RoundedBorder()).BorderForeground(surface0).
Render(contents)
}
+6 -164
View File
@@ -151,171 +151,13 @@ func (s *ViewerTestSuite) TestViewerUpdate() {
require.False(m.SearchBar.TextInput.Focused(), "expected search bar to be unfocused, got focused")
// quit the program with 'q'
_, cmd := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("q")})
fmt.Printf("cmd: %v\n", cmd)
require.Equal(tea.Quit(), cmd(), "expected quit command, got %v", cmd)
_, command := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("q")})
fmt.Printf("cmd: %v\n", command)
require.Equal(tea.Quit(), command(), "expected quit command, got %v", command)
// quit the program with ctrl+c
_, cmd = m.Update(tea.KeyMsg{Type: tea.KeyCtrlC})
fmt.Printf("cmd: %v\n", cmd)
require.Equal(tea.Quit(), cmd(), "expected quit command, got %v", cmd)
}
func (s *ViewerTestSuite) TestListScrolling() {
t := s.T()
// create new ui model
m, err := viewer.NewModel(s.maxTimestamp, s.minTimestamp, s.useCurrentTime, s.db)
require.NoError(t, err)
// get current selected index
initialSelectedIndex := m.List.Rows.Index()
// use down key to scroll the list down five times
for i := 0; i < 5; i++ {
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyDown,
},
))
}
// verify that the list was scrolled down five times from the initially selected index
require.Equal(t, initialSelectedIndex+5, m.List.Rows.Index())
// use up key to scroll the list up three times
for i := 0; i < 3; i++ {
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyUp,
},
))
}
// verify that the list was scrolled up 3 times, resulting in an index of 2 away from the initial index
require.Equal(t, initialSelectedIndex+2, m.List.Rows.Index())
}
func (s *ViewerTestSuite) TestListPaging() {
t := s.T()
// create new ui model
m, err := viewer.NewModel(s.maxTimestamp, s.minTimestamp, s.useCurrentTime, s.db)
require.NoError(t, err)
// get current page
initialPage := m.List.Rows.Paginator.Page
// select the 5th row in the list to ensure that the cursor is kept on the same row while paging
cursor := 5
m.List.Rows.Select(cursor)
// get current selected index
initialSelectedIndex := m.List.Rows.Index()
// use page down key to page down 5 pages in the list
for i := 0; i < 5; i++ {
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyPgDown,
},
))
}
// verify that the list was paged down five times from the initial page
require.Equal(t, initialPage+5, m.List.Rows.Paginator.Page, "after paging down 5 times, expected page to be %d, got %d", initialPage+5, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly
require.Equal(t, initialSelectedIndex+(m.List.Rows.Paginator.PerPage*5), m.List.Rows.Index(), "after paging down 5 times, expected selected index to be %d, got %d", initialSelectedIndex+(m.List.Rows.Paginator.PerPage*5), m.List.Rows.Index())
// verify that the cursor is still on the 5th row
require.Equal(t, cursor, m.List.Rows.Cursor(), "after paging down 5 times, expected cursor to remain as %d, got %d", cursor, m.List.Rows.Cursor())
// use page up key to scroll the list up 3 pages
for i := 0; i < 3; i++ {
m.Update(
tea.KeyMsg{
Type: tea.KeyPgUp,
},
)
}
// verify that the list was paged up 3 times, resulting in 2 pages away from the initial page (since we paged down 5 times first)
require.Equal(t, initialPage+2, m.List.Rows.Paginator.Page, "after paging up 3 times, expected page to be %d, got %d", initialPage+2, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly
require.Equal(t, initialSelectedIndex+(m.List.Rows.Paginator.PerPage*2), m.List.Rows.Index(), "after paging up 3 times, expected index to be %d, got %d", initialSelectedIndex+(m.List.Rows.Paginator.PerPage*2), m.List.Rows.Index())
// verify that the cursor is still on the correct row
require.Equal(t, cursor, m.List.Rows.Cursor(), "after paging up 3 times, expected cursor to remain as %d, got %d", cursor, m.List.Rows.Cursor())
// use home key to scroll to the start of the list
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyHome,
},
))
// verify that the list was paged to the start
require.Equal(t, 0, m.List.Rows.Paginator.Page, "after paging to the start, expected page to be %d, got %d", 0, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly
require.Equal(t, initialSelectedIndex, m.List.Rows.Index(), "after paging to the start, expected index to be %d, got %d", initialSelectedIndex, m.List.Rows.Index())
// verify that the cursor is still on the correct row
require.Equal(t, cursor, m.List.Rows.Cursor(), "after paging to the start, expected cursor to remain as %d, got %d", cursor, m.List.Rows.Cursor())
// use page end key to scroll to the end of the list
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyEnd,
},
))
// verify that the list was paged to the end
require.Equal(t, m.List.Rows.Paginator.TotalPages-1, m.List.Rows.Paginator.Page, "after paging to the end, expected page to be %d, got %d", m.List.Rows.Paginator.TotalPages-1, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly (since the last page may have fewer items than the cursor index, the selected index should be min(cursor, items on last page - 1 ))
endCursor := min(cursor, m.List.Rows.Paginator.ItemsOnPage(len(m.List.Rows.Items()))-1)
endIndex := (m.List.Rows.Paginator.Page * m.List.Rows.Paginator.PerPage) + endCursor
require.Equal(t, endIndex, m.List.Rows.Index(), "after paging to the end, expected selected index to be %d, got %d", endIndex, m.List.Rows.Index())
// verify that the cursor was updated to the correct spot
require.Equal(t, endCursor, m.List.Rows.Cursor(), "after paging to the end, expected cursor to be %d, got %d", endCursor, m.List.Rows.Cursor())
// page up one page (to the second-to-last page)
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyPgUp,
},
))
// set the cursor down to the bottom row of the second-to-last page to ensure that the selected row is greater than the number of items on the last page we will page to
m.List.Rows.Select((m.List.Rows.Paginator.Page * m.List.Rows.Paginator.PerPage) + (m.List.Rows.Paginator.ItemsOnPage(len(m.List.Rows.Items())) - 1))
// page down one page (back to last page)
m.Update(tea.KeyMsg(
tea.Key{
Type: tea.KeyPgDown,
},
))
// // verify that the list was paged up to the last page
require.Equal(t, m.List.Rows.Paginator.TotalPages-1, m.List.Rows.Paginator.Page, "after paging up to the last page, expected page to be %d, got %d", m.List.Rows.Paginator.TotalPages-1, m.List.Rows.Paginator.Page)
// verify that the selected index was updated accordingly to the last item on the last page
require.Equal(t, endIndex, m.List.Rows.Index(), "after paging up to the last page, expected selected index to be %d, got %d", endIndex, m.List.Rows.Index())
// verify that the cursor was updated to the correct spot
require.Equal(t, endCursor, m.List.Rows.Cursor(), "after paging up to the last page, expected cursor to be %d, got %d", endCursor, m.List.Rows.Cursor())
// fmt.Printf("page: %d\n", m.List.Rows.Paginator.Page)
// fmt.Printf("initialSelectedIndex: %d\n", initialSelectedIndex)
// fmt.Printf("actual index: %d\n", m.List.Rows.Index())
// fmt.Printf("perPage: %d\n", m.List.Rows.Paginator.PerPage)
// fmt.Printf("items on page: %d\n", m.List.Rows.Paginator.ItemsOnPage(len(m.List.Rows.Items())))
// fmt.Printf("cursor: %d\n", m.List.Rows.Cursor())
_, command = m.Update(tea.KeyMsg{Type: tea.KeyCtrlC})
fmt.Printf("cmd: %v\n", command)
require.Equal(tea.Quit(), command(), "expected quit command, got %v", command)
}