From 15f39b6ea5e6ea2e54d3e0b65ad6133628e7af16 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=B4mulo=20Farias?= <romulo.farias@elastic.co>
Date: Mon, 25 Nov 2024 15:24:36 +0100
Subject: [PATCH] Add entities.metadata

---
 internal/inventory/asset.go          | 23 +++++++++++++++++++++++
 internal/inventory/inventory.go      |  1 +
 internal/inventory/inventory_test.go |  8 ++++++++
 3 files changed, 32 insertions(+)

diff --git a/internal/inventory/asset.go b/internal/inventory/asset.go
index 33f85a2f89..431f3ed273 100644
--- a/internal/inventory/asset.go
+++ b/internal/inventory/asset.go
@@ -329,9 +329,32 @@ type AssetResourcePolicy struct {
 	Condition  map[string]any `json:"condition,omitempty"`
 }
 
+// EntityMetadata maps metadata required to use Entity Store
+type entityMetadata struct {
+	Category string `json:"category"`
+	Type     string `json:"type"`
+}
+
 // AssetEnricher functional builder function
 type AssetEnricher func(asset *AssetEvent)
 
+func (a *AssetEvent) getEntityMetadata() map[string]entityMetadata {
+	ids := a.Asset.Id
+
+	if len(ids) == 0 {
+		return nil
+	}
+
+	// Picking up only first id, we need to make a decision on if we
+	// have a "primary" id or if we duplicate data
+	return map[string]entityMetadata{
+		ids[0]: {
+			Category: string(a.Asset.Category),
+			Type:     string(a.Asset.Type),
+		},
+	}
+}
+
 func NewAssetEvent(c AssetClassification, ids []string, name string, enrichers ...AssetEnricher) AssetEvent {
 	a := AssetEvent{
 		Asset: Asset{
diff --git a/internal/inventory/inventory.go b/internal/inventory/inventory.go
index c5a5b171ca..83c2c73219 100644
--- a/internal/inventory/inventory.go
+++ b/internal/inventory/inventory.go
@@ -114,6 +114,7 @@ func (a *AssetInventory) publish(assets []AssetEvent) {
 				"iam":               e.IAM,
 				"resource_policies": e.ResourcePolicies,
 				"related.entity":    e.Asset.Id,
+				"entities":          map[string]any{"metadata": e.getEntityMetadata()},
 			},
 		}
 	})
diff --git a/internal/inventory/inventory_test.go b/internal/inventory/inventory_test.go
index 685d951634..87dad4f048 100644
--- a/internal/inventory/inventory_test.go
+++ b/internal/inventory/inventory_test.go
@@ -86,6 +86,14 @@ func TestAssetInventory_Run(t *testing.T) {
 					},
 				},
 				"related.entity": []string{"arn:aws:ec2:us-east::ec2/234567890"},
+				"entities": map[string]any{
+					"metadata": map[string]entityMetadata{
+						"arn:aws:ec2:us-east::ec2/234567890": {
+							Category: string(CategoryInfrastructure),
+							Type:     string(TypeVirtualMachine),
+						},
+					},
+				},
 			},
 		},
 	}