mirror of
https://github.com/activecm/rita
synced 2026-06-08 13:02:45 +00:00
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:
+70
-10
@@ -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
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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, `
|
||||
|
||||
@@ -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")
|
||||
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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(), `
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
@@ -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()),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user