pkg/cticlient: Add missing field in SmokeItem and FireItem (#3413)

* pkg/cticlient: Add missing field in SmokeItem and FireItem
crowdsecurity · Jan 20, 2025 · d6b3841 · d6b3841
1 parent ff9b350
commit d6b3841
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 2 deletions.
diff --git a/pkg/cticlient/types.go b/pkg/cticlient/types.go
@@ -64,6 +64,9 @@ type CTIReferences struct {
 type SmokeItem struct {
 	IpRangeScore         int                 `json:"ip_range_score"`
 	Ip                   string              `json:"ip"`
+	Reputation           string              `json:"reputation"`
+	BackgroundNoise      string              `json:"background_noise"`
+	Confidence           string              `json:"confidence"`
 	IpRange              *string             `json:"ip_range"`
 	AsName               *string             `json:"as_name"`
 	AsNum                *int                `json:"as_num"`
@@ -77,6 +80,7 @@ type SmokeItem struct {
 	BackgroundNoiseScore *int                `json:"background_noise_score"`
 	Scores               CTIScores           `json:"scores"`
 	References           []CTIReferences     `json:"references"`
+	CVEs                 []string            `json:"cves"`
 	IsOk                 bool                `json:"-"`
 }
 
@@ -120,6 +124,10 @@ type FireItem struct {
 	BackgroundNoiseScore *int                `json:"background_noise_score"`
 	Scores               CTIScores           `json:"scores"`
 	References           []CTIReferences     `json:"references"`
+	CVEs                 []string            `json:"cves"`
+	Reputation           string              `json:"reputation"`
+	BackgroundNoise      string              `json:"background_noise"`
+	Confidence           string              `json:"confidence"`
 	State                string              `json:"state"`
 	Expiration           CustomTime          `json:"expiration"`
 }
@@ -209,6 +217,18 @@ func (c *SmokeItem) GetFalsePositives() []string {
 	return ret
 }
 
+func (c *SmokeItem) GetClassifications() []string {
+	ret := make([]string, 0)
+
+	if c.Classifications.Classifications != nil {
+		for _, b := range c.Classifications.Classifications {
+			ret = append(ret, b.Name)
+		}
+	}
+
+	return ret
+}
+
 func (c *SmokeItem) IsFalsePositive() bool {
 	if c.Classifications.FalsePositives != nil {
 		if len(c.Classifications.FalsePositives) > 0 {
@@ -282,6 +302,18 @@ func (c *FireItem) GetFalsePositives() []string {
 	return ret
 }
 
+func (c *FireItem) GetClassifications() []string {
+	ret := make([]string, 0)
+
+	if c.Classifications.Classifications != nil {
+		for _, b := range c.Classifications.Classifications {
+			ret = append(ret, b.Name)
+		}
+	}
+
+	return ret
+}
+
 func (c *FireItem) IsFalsePositive() bool {
 	if c.Classifications.FalsePositives != nil {
 		if len(c.Classifications.FalsePositives) > 0 {

diff --git a/pkg/cticlient/types_test.go b/pkg/cticlient/types_test.go
@@ -40,8 +40,14 @@ func getSampleSmokeItem() SmokeItem {
 			DaysAge:   1,
 		},
 		Classifications: CTIClassifications{
-			FalsePositives:  []CTIClassification{},
-			Classifications: []CTIClassification{},
+			FalsePositives: []CTIClassification{},
+			Classifications: []CTIClassification{
+				{
+					Name:        "profile:likely_botnet",
+					Label:       "Likely Botnet",
+					Description: "IP appears to be a botnet.",
+				},
+			},
 		},
 		AttackDetails: []*CTIAttackDetails{
 			{
@@ -101,6 +107,7 @@ func TestBasicSmokeItem(t *testing.T) {
 	assert.Equal(t, 3, item.GetBackgroundNoiseScore())
 	assert.Equal(t, []string{}, item.GetFalsePositives())
 	assert.False(t, item.IsFalsePositive())
+	assert.Equal(t, []string{"profile:likely_botnet"}, item.GetClassifications())
 }
 
 func TestEmptySmokeItem(t *testing.T) {
@@ -112,4 +119,5 @@ func TestEmptySmokeItem(t *testing.T) {
 	assert.Equal(t, 0, item.GetBackgroundNoiseScore())
 	assert.Equal(t, []string{}, item.GetFalsePositives())
 	assert.False(t, item.IsFalsePositive())
+	assert.Equal(t, []string{}, item.GetClassifications())
 }