GetActiveDecisionsCount()and GetActiveDecisionsTimeLeft()expr helpers (#3013)

This commit is contained in:
blotus 2024-05-15 15:33:43 +02:00 committed by GitHub
parent cc63729b2c
commit 1b894a292b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 367 additions and 1 deletions

View file

@ -640,6 +640,68 @@ func (c *Client) CountDecisionsByValue(decisionValue string) (int, error) {
return count, nil
}
func (c *Client) CountActiveDecisionsByValue(decisionValue string) (int, error) {
var err error
var start_ip, start_sfx, end_ip, end_sfx int64
var ip_sz, count int
ip_sz, start_ip, start_sfx, end_ip, end_sfx, err = types.Addr2Ints(decisionValue)
if err != nil {
return 0, fmt.Errorf("unable to convert '%s' to int: %s", decisionValue, err)
}
contains := true
decisions := c.Ent.Decision.Query()
decisions, err = applyStartIpEndIpFilter(decisions, contains, ip_sz, start_ip, start_sfx, end_ip, end_sfx)
if err != nil {
return 0, fmt.Errorf("fail to apply StartIpEndIpFilter: %w", err)
}
decisions = decisions.Where(decision.UntilGT(time.Now().UTC()))
count, err = decisions.Count(c.CTX)
if err != nil {
return 0, fmt.Errorf("fail to count decisions: %w", err)
}
return count, nil
}
func (c *Client) GetActiveDecisionsTimeLeftByValue(decisionValue string) (time.Duration, error) {
var err error
var start_ip, start_sfx, end_ip, end_sfx int64
var ip_sz int
ip_sz, start_ip, start_sfx, end_ip, end_sfx, err = types.Addr2Ints(decisionValue)
if err != nil {
return 0, fmt.Errorf("unable to convert '%s' to int: %s", decisionValue, err)
}
contains := true
decisions := c.Ent.Decision.Query().Where(
decision.UntilGT(time.Now().UTC()),
)
decisions, err = applyStartIpEndIpFilter(decisions, contains, ip_sz, start_ip, start_sfx, end_ip, end_sfx)
if err != nil {
return 0, fmt.Errorf("fail to apply StartIpEndIpFilter: %w", err)
}
decisions = decisions.Order(ent.Desc(decision.FieldUntil))
decision, err := decisions.First(c.CTX)
if err != nil && !ent.IsNotFound(err) {
return 0, fmt.Errorf("fail to get decision: %w", err)
}
if decision == nil {
return 0, nil
}
return decision.Until.Sub(time.Now().UTC()), nil
}
func (c *Client) CountDecisionsSinceByValue(decisionValue string, since time.Time) (int, error) {
ip_sz, start_ip, start_sfx, end_ip, end_sfx, err := types.Addr2Ints(decisionValue)

View file

@ -231,6 +231,20 @@ var exprFuncs = []exprCustomFunc{
new(func(string) int),
},
},
{
name: "GetActiveDecisionsCount",
function: GetActiveDecisionsCount,
signature: []interface{}{
new(func(string) int),
},
},
{
name: "GetActiveDecisionsTimeLeft",
function: GetActiveDecisionsTimeLeft,
signature: []interface{}{
new(func(string) time.Duration),
},
},
{
name: "GetDecisionsSinceCount",
function: GetDecisionsSinceCount,

View file

@ -1118,6 +1118,268 @@ func TestGetDecisionsSinceCount(t *testing.T) {
}
}
func TestGetActiveDecisionsCount(t *testing.T) {
existingIP := "1.2.3.4"
unknownIP := "1.2.3.5"
ip_sz, start_ip, start_sfx, end_ip, end_sfx, err := types.Addr2Ints(existingIP)
if err != nil {
t.Errorf("unable to convert '%s' to int: %s", existingIP, err)
}
// Add sample data to DB
dbClient = getDBClient(t)
decision := dbClient.Ent.Decision.Create().
SetUntil(time.Now().UTC().Add(time.Hour)).
SetScenario("crowdsec/test").
SetStartIP(start_ip).
SetStartSuffix(start_sfx).
SetEndIP(end_ip).
SetEndSuffix(end_sfx).
SetIPSize(int64(ip_sz)).
SetType("ban").
SetScope("IP").
SetValue(existingIP).
SetOrigin("CAPI").
SaveX(context.Background())
if decision == nil {
require.Error(t, errors.Errorf("Failed to create sample decision"))
}
expiredDecision := dbClient.Ent.Decision.Create().
SetUntil(time.Now().UTC().Add(-time.Hour)).
SetScenario("crowdsec/test").
SetStartIP(start_ip).
SetStartSuffix(start_sfx).
SetEndIP(end_ip).
SetEndSuffix(end_sfx).
SetIPSize(int64(ip_sz)).
SetType("ban").
SetScope("IP").
SetValue(existingIP).
SetOrigin("CAPI").
SaveX(context.Background())
if expiredDecision == nil {
require.Error(t, errors.Errorf("Failed to create sample decision"))
}
err = Init(dbClient)
require.NoError(t, err)
tests := []struct {
name string
env map[string]interface{}
code string
result string
err string
}{
{
name: "GetActiveDecisionsCount() test: existing IP count",
env: map[string]interface{}{
"Alert": &models.Alert{
Source: &models.Source{
Value: &existingIP,
},
Decisions: []*models.Decision{
{
Value: &existingIP,
},
},
},
},
code: "Sprintf('%d', GetActiveDecisionsCount(Alert.GetValue()))",
result: "1",
err: "",
},
{
name: "GetActiveDecisionsCount() test: unknown IP count",
env: map[string]interface{}{
"Alert": &models.Alert{
Source: &models.Source{
Value: &unknownIP,
},
Decisions: []*models.Decision{
{
Value: &unknownIP,
},
},
},
},
code: "Sprintf('%d', GetActiveDecisionsCount(Alert.GetValue()))",
result: "0",
err: "",
},
}
for _, test := range tests {
program, err := expr.Compile(test.code, GetExprOptions(test.env)...)
require.NoError(t, err)
output, err := expr.Run(program, test.env)
require.NoError(t, err)
require.Equal(t, test.result, output)
log.Printf("test '%s' : OK", test.name)
}
}
func TestGetActiveDecisionsTimeLeft(t *testing.T) {
existingIP := "1.2.3.4"
unknownIP := "1.2.3.5"
ip_sz, start_ip, start_sfx, end_ip, end_sfx, err := types.Addr2Ints(existingIP)
if err != nil {
t.Errorf("unable to convert '%s' to int: %s", existingIP, err)
}
// Add sample data to DB
dbClient = getDBClient(t)
decision := dbClient.Ent.Decision.Create().
SetUntil(time.Now().UTC().Add(time.Hour)).
SetScenario("crowdsec/test").
SetStartIP(start_ip).
SetStartSuffix(start_sfx).
SetEndIP(end_ip).
SetEndSuffix(end_sfx).
SetIPSize(int64(ip_sz)).
SetType("ban").
SetScope("IP").
SetValue(existingIP).
SetOrigin("CAPI").
SaveX(context.Background())
if decision == nil {
require.Error(t, errors.Errorf("Failed to create sample decision"))
}
longerDecision := dbClient.Ent.Decision.Create().
SetUntil(time.Now().UTC().Add(2 * time.Hour)).
SetScenario("crowdsec/test").
SetStartIP(start_ip).
SetStartSuffix(start_sfx).
SetEndIP(end_ip).
SetEndSuffix(end_sfx).
SetIPSize(int64(ip_sz)).
SetType("ban").
SetScope("IP").
SetValue(existingIP).
SetOrigin("CAPI").
SaveX(context.Background())
if longerDecision == nil {
require.Error(t, errors.Errorf("Failed to create sample decision"))
}
err = Init(dbClient)
require.NoError(t, err)
tests := []struct {
name string
env map[string]interface{}
code string
min float64
max float64
err string
}{
{
name: "GetActiveDecisionsTimeLeft() test: existing IP time left",
env: map[string]interface{}{
"Alert": &models.Alert{
Source: &models.Source{
Value: &existingIP,
},
Decisions: []*models.Decision{
{
Value: &existingIP,
},
},
},
},
code: "GetActiveDecisionsTimeLeft(Alert.GetValue())",
min: 7195, // 5 seconds margin to make sure the test doesn't fail randomly in the CI
max: 7200,
err: "",
},
{
name: "GetActiveDecisionsTimeLeft() test: unknown IP time left",
env: map[string]interface{}{
"Alert": &models.Alert{
Source: &models.Source{
Value: &unknownIP,
},
Decisions: []*models.Decision{
{
Value: &unknownIP,
},
},
},
},
code: "GetActiveDecisionsTimeLeft(Alert.GetValue())",
min: 0,
max: 0,
err: "",
},
{
name: "GetActiveDecisionsTimeLeft() test: existing IP and call time.Duration method",
env: map[string]interface{}{
"Alert": &models.Alert{
Source: &models.Source{
Value: &existingIP,
},
Decisions: []*models.Decision{
{
Value: &existingIP,
},
},
},
},
code: "GetActiveDecisionsTimeLeft(Alert.GetValue()).Hours()",
min: 2,
max: 2,
},
{
name: "GetActiveDecisionsTimeLeft() test: unknown IP and call time.Duration method",
env: map[string]interface{}{
"Alert": &models.Alert{
Source: &models.Source{
Value: &unknownIP,
},
Decisions: []*models.Decision{
{
Value: &unknownIP,
},
},
},
},
code: "GetActiveDecisionsTimeLeft(Alert.GetValue()).Hours()",
min: 0,
max: 0,
},
}
delta := 0.0001
for _, test := range tests {
program, err := expr.Compile(test.code, GetExprOptions(test.env)...)
require.NoError(t, err)
output, err := expr.Run(program, test.env)
require.NoError(t, err)
switch o := output.(type) {
case time.Duration:
require.LessOrEqual(t, int(o.Seconds()), int(test.max))
require.GreaterOrEqual(t, int(o.Seconds()), int(test.min))
case float64:
require.LessOrEqual(t, o, test.max+delta)
require.GreaterOrEqual(t, o, test.min-delta)
default:
t.Fatalf("GetActiveDecisionsTimeLeft() should return a time.Duration or a float64")
}
}
}
func TestParseUnixTime(t *testing.T) {
tests := []struct {
name string

View file

@ -550,7 +550,7 @@ func GetDecisionsSinceCount(params ...any) (any, error) {
value := params[0].(string)
since := params[1].(string)
if dbClient == nil {
log.Error("No database config to call GetDecisionsCount()")
log.Error("No database config to call GetDecisionsSinceCount()")
return 0, nil
}
sinceDuration, err := time.ParseDuration(since)
@ -567,6 +567,34 @@ func GetDecisionsSinceCount(params ...any) (any, error) {
return count, nil
}
func GetActiveDecisionsCount(params ...any) (any, error) {
value := params[0].(string)
if dbClient == nil {
log.Error("No database config to call GetActiveDecisionsCount()")
return 0, nil
}
count, err := dbClient.CountActiveDecisionsByValue(value)
if err != nil {
log.Errorf("Failed to get active decisions count from value '%s'", value)
return 0, err
}
return count, nil
}
func GetActiveDecisionsTimeLeft(params ...any) (any, error) {
value := params[0].(string)
if dbClient == nil {
log.Error("No database config to call GetActiveDecisionsTimeLeft()")
return 0, nil
}
timeLeft, err := dbClient.GetActiveDecisionsTimeLeftByValue(value)
if err != nil {
log.Errorf("Failed to get active decisions time left from value '%s'", value)
return 0, err
}
return timeLeft, nil
}
// func LookupHost(value string) []string {
func LookupHost(params ...any) (any, error) {
value := params[0].(string)