OLED-246 - refactor and add eks enricher for account id and account g…

…roup (#10)

* refactor and add eks enricher for account_id

* change account id field name

* WIP

* add env var

* remove testing lib

* remove aws from account group

* change account group to canva function

* change name

enricher/ecs/ecs.go

@@ -0,0 +1,102 @@
+package ecs
+
+import (
+	"os"
+	"regexp"
+	"strconv"
+	"time"
+
+	"github.com/canva/amazon-kinesis-streams-for-fluent-bit/enricher"
+	"github.com/canva/amazon-kinesis-streams-for-fluent-bit/enricher/mappings"
+	"github.com/sirupsen/logrus"
+)
+
+type Enricher struct {
+	canvaAWSAccount string
+	canvaAppName    string
+	logGroup        string
+	ecsTaskFamily   string
+	ecsTaskRevision int
+}
+
+var _ enricher.IEnricher = (*Enricher)(nil)
+
+func NewEnricher() *Enricher {
+	ecsTaskDefinition := os.Getenv("ECS_TASK_DEFINITION")
+	re := regexp.MustCompile(`^(?P<ecs_task_family>[^ ]*):(?P<ecs_task_revision>[\d]+)$`)
+	ecsTaskDefinitionParts := re.FindStringSubmatch(ecsTaskDefinition)
+	var (
+		ecsTaskFamily   string
+		ecsTaskRevision int
+	)
+	ecsTaskFamilyIndex := re.SubexpIndex("ecs_task_family")
+	ecsTaskRevisionIndex := re.SubexpIndex("ecs_task_revision")
+
+	if len(ecsTaskDefinitionParts) >= ecsTaskFamilyIndex {
+		ecsTaskFamily = ecsTaskDefinitionParts[ecsTaskFamilyIndex]
+	}
+	if len(ecsTaskDefinitionParts) >= ecsTaskRevisionIndex {
+		var err error
+		ecsTaskRevision, err = strconv.Atoi(ecsTaskDefinitionParts[re.SubexpIndex("ecs_task_revision")])
+		if err != nil {
+			logrus.Warnf("[kinesis] ecs_task_revision not found for ECS_TASK_DEFINITION=%s", ecsTaskDefinition)
+		}
+	}
+
+	return &Enricher{
+		canvaAWSAccount: os.Getenv("CANVA_AWS_ACCOUNT"),
+		canvaAppName:    os.Getenv("CANVA_APP_NAME"),
+		logGroup:        os.Getenv("LOG_GROUP"),
+		ecsTaskFamily:   ecsTaskFamily,
+		ecsTaskRevision: ecsTaskRevision,
+	}
+}
+
+// EnrichRecord modifies existing record.
+func (enr *Enricher) EnrichRecord(r map[interface{}]interface{}, t time.Time) map[interface{}]interface{} {
+	resource := map[interface{}]interface{}{
+		mappings.RESOURCE_CLOUD_ACCOUNT_ID: enr.canvaAWSAccount,
+		"service.name":                     enr.canvaAppName,
+		"cloud.platform":                   "aws_ecs",
+		"aws.ecs.launchtype":               "EC2",
+		"aws.ecs.task.family":              enr.ecsTaskFamily,
+		"aws.ecs.task.revision":            enr.ecsTaskRevision,
+		"aws.log.group.names":              enr.logGroup,
+	}
+	body := make(map[interface{}]interface{})
+
+	var (
+		ok        bool
+		strVal    string
+		timestamp interface{}
+	)
+	for k, v := range r {
+		strVal, ok = k.(string)
+		if ok {
+			switch strVal {
+			case "ecs_task_definition":
+				// Skip
+			case "timestamp":
+				timestamp = v
+			case "ec2_instance_id":
+				resource["host.id"] = v
+			case "ecs_cluster":
+				resource["aws.ecs.cluster.name"] = v
+			case "ecs_task_arn":
+				resource["aws.ecs.task.arn"] = v
+			case "container_id":
+				resource["container.id"] = v
+			case "container_name":
+				resource["container.name"] = v
+			default:
+				body[k] = v
+			}
+		}
+	}
+	return map[interface{}]interface{}{
+		"resource":          resource,
+		"body":              body,
+		"timestamp":         timestamp,
+		"observedTimestamp": t.UnixMilli(),
+	}
+}

enricher/ecs/ecs_test.go

@@ -0,0 +1,80 @@
+package ecs
+
+import (
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/canva/amazon-kinesis-streams-for-fluent-bit/enricher"
+	"github.com/canva/amazon-kinesis-streams-for-fluent-bit/enricher/mappings"
+)
+
+func TestEnrichRecords(t *testing.T) {
+	type args struct {
+		r map[interface{}]interface{}
+		t time.Time
+	}
+	tests := []struct {
+		name string
+		enr  enricher.IEnricher
+		args args
+		want map[interface{}]interface{}
+	}{
+		{
+			name: "enrich",
+			enr: &Enricher{
+				canvaAWSAccount: "canva_aws_account_val",
+				canvaAppName:    "canva_app_name_val",
+				logGroup:        "log_group_val",
+				ecsTaskFamily:   "ecs_task_family_val",
+				ecsTaskRevision: 10001,
+			},
+			args: args{
+				map[interface{}]interface{}{
+					"ec2_instance_id":     "ec2_instance_id_val",
+					"ecs_cluster":         "ecs_cluster_val",
+					"ecs_task_arn":        "ecs_task_arn_val",
+					"container_id":        "container_id_val",
+					"container_name":      "container_name_val",
+					"other_key_1":         "other_value_1",
+					"other_key_2":         "other_value_2",
+					"other_key_3":         "other_value_3",
+					"timestamp":           "1234567890",
+					"ecs_task_definition": "ecs_task_definition_val",
+				},
+				time.Date(2009, time.November, 10, 23, 7, 5, 432000000, time.UTC),
+			},
+			want: map[interface{}]interface{}{
+				"resource": map[interface{}]interface{}{
+					mappings.RESOURCE_CLOUD_ACCOUNT_ID: "canva_aws_account_val",
+					"service.name":                     "canva_app_name_val",
+					"cloud.platform":                   "aws_ecs",
+					"aws.ecs.launchtype":               "EC2",
+					"aws.ecs.task.family":              "ecs_task_family_val",
+					"aws.ecs.task.revision":            10001,
+					"aws.log.group.names":              "log_group_val",
+					"host.id":                          "ec2_instance_id_val",
+					"aws.ecs.cluster.name":             "ecs_cluster_val",
+					"aws.ecs.task.arn":                 "ecs_task_arn_val",
+					"container.id":                     "container_id_val",
+					"container.name":                   "container_name_val",
+				},
+				"body": map[interface{}]interface{}{
+					"other_key_1": "other_value_1",
+					"other_key_2": "other_value_2",
+					"other_key_3": "other_value_3",
+				},
+				"timestamp":         "1234567890",
+				"observedTimestamp": int64(1257894425432),
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := tt.enr.EnrichRecord(tt.args.r, tt.args.t)
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("enricher.enrichRecord() = %+v, want %+v", got, tt.want)
+			}
+		})
+	}
+}

enricher/eks/eks.go

@@ -0,0 +1,37 @@
+package eks
+
+import (
+	"time"
+
+	"github.com/caarlos0/env/v7"
+	"github.com/canva/amazon-kinesis-streams-for-fluent-bit/enricher"
+	"github.com/canva/amazon-kinesis-streams-for-fluent-bit/enricher/mappings"
+)
+
+type Enricher struct {
+	// AWS Account ID
+	AccountId string `env:"CANVA_AWS_ACCOUNT,required"`
+	// Canva Account Group Function
+	CanvaAccountFunction string `env:"CANVA_ACCOUNT_FUNCTION,required"`
+}
+
+func NewEnricher() (*Enricher, error) {
+	enricher := Enricher{}
+	if err := env.Parse(&enricher); err != nil {
+		return nil, err
+	}
+
+	return &enricher, nil
+}
+
+var _ enricher.IEnricher = (*Enricher)(nil)
+
+func (e Enricher) EnrichRecord(r map[interface{}]interface{}, _ time.Time) map[interface{}]interface{} {
+	// add resource attributes
+	r["resource"] = map[interface{}]interface{}{
+		mappings.RESOURCE_CLOUD_ACCOUNT_ID: e.AccountId,
+		mappings.RESOURCE_ACCOUNT_GROUP:    e.CanvaAccountFunction,
+	}
+
+	return r
+}

enricher/eks/eks_test.go

@@ -0,0 +1,119 @@
+package eks
+
+import (
+	"testing"
+	"time"
+
+	"github.com/canva/amazon-kinesis-streams-for-fluent-bit/enricher/mappings"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestValidNewEnricher(t *testing.T) {
+	var cases = []struct {
+		Name     string
+		Env      map[string]string
+		Expected *Enricher
+	}{
+		{
+			Name: "Gets AccountId",
+			Env: map[string]string{
+				mappings.ENV_ACCOUNT_ID:    "1234567890",
+				mappings.ENV_ACCOUNT_GROUP: DummyAccountGroup,
+			},
+			Expected: &Enricher{
+				AccountId:            "1234567890",
+				CanvaAccountFunction: DummyAccountGroup,
+			},
+		},
+		{
+			Name: "Gets Account Group",
+			Env: map[string]string{
+				mappings.ENV_ACCOUNT_ID:    DummyAccountId,
+				mappings.ENV_ACCOUNT_GROUP: "PII",
+			},
+			Expected: &Enricher{
+				AccountId:            DummyAccountId,
+				CanvaAccountFunction: "PII",
+			},
+		},
+	}
+
+	for _, v := range cases {
+		t.Run(v.Name, func(tt *testing.T) {
+			for k, v := range v.Env {
+				tt.Setenv(k, v)
+			}
+			actual, err := NewEnricher()
+
+			assert.NoError(tt, err)
+
+			assert.Equal(tt, v.Expected, actual)
+
+			tt.Cleanup(func() {})
+		})
+	}
+}
+
+func TestInvalidNewEnricher(t *testing.T) {
+	enricher, err := NewEnricher()
+
+	assert.Nil(t, enricher)
+	assert.Error(t, err)
+}
+
+func TestEnrichRecordsWithAccountId(t *testing.T) {
+	var cases = []struct {
+		Name     string
+		Enricher Enricher
+		Input    map[interface{}]interface{}
+		Expected map[interface{}]interface{}
+	}{
+		{
+			Name: "Adds Account Id",
+			Enricher: Enricher{
+				AccountId:            "1234567",
+				CanvaAccountFunction: DummyAccountGroup,
+			},
+			Input: map[interface{}]interface{}{
+				"log": "hello world",
+			},
+			Expected: map[interface{}]interface{}{
+				"log": "hello world",
+				"resource": map[interface{}]interface{}{
+					mappings.RESOURCE_CLOUD_ACCOUNT_ID: "1234567",
+					mappings.RESOURCE_ACCOUNT_GROUP:    DummyAccountGroup,
+				},
+			},
+		},
+		{
+			Name: "Adds Account Group",
+			Enricher: Enricher{
+				AccountId:            DummyAccountId,
+				CanvaAccountFunction: "PII",
+			},
+			Input: map[interface{}]interface{}{
+				"log": "hello world",
+			},
+			Expected: map[interface{}]interface{}{
+				"log": "hello world",
+				"resource": map[interface{}]interface{}{
+					mappings.RESOURCE_CLOUD_ACCOUNT_ID: DummyAccountId,
+					mappings.RESOURCE_ACCOUNT_GROUP:    "PII",
+				},
+			},
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.Name, func(tt *testing.T) {
+			actual := c.Enricher.EnrichRecord(c.Input, DummyTime)
+			assert.Equal(tt, c.Expected, actual)
+		})
+	}
+}
+
+var (
+	DummyTime         = time.Now()
+	DummyAccountGroup = "general"
+	DummyAccountId    = "Account_Id"
+)