From 3cbb651bd191c1e2940a81fc4491da11aae208f1 Mon Sep 17 00:00:00 2001
From: Youssef <53960652+yti93@users.noreply.github.com>
Date: Wed, 5 Aug 2026 13:30:31 -0700
Subject: [PATCH] fix(wiz): derive OpenSearch region from ARN

Amp-Thread-ID: https://ampcode.com/threads/T-019fd393-6ce9-744e-8b29-d0165a10d3c4
Co-authored-by: Amp <amp@ampcode.com>
---
 pkg/inventory/wiz/generic.go      | 24 ++++++++++-
 pkg/inventory/wiz/generic_test.go | 70 +++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/pkg/inventory/wiz/generic.go b/pkg/inventory/wiz/generic.go
index 4f9a240..7c98dfd 100644
--- a/pkg/inventory/wiz/generic.go
+++ b/pkg/inventory/wiz/generic.go
@@ -5,9 +5,11 @@ import (
 	"encoding/json"
 	"log/slog"
 	"os"
+	"regexp"
 	"strings"
 	"time"
 
+	"github.com/aws/aws-sdk-go-v2/aws/arn"
 	"github.com/pkg/errors"
 
 	"github.com/block/Version-Guard/pkg/config"
@@ -135,6 +137,8 @@ var wellKnownFieldMappingKeys = map[string]struct{}{
 	"tags":        {},
 }
 
+var awsRegionPattern = regexp.MustCompile(`^[a-z]{2}(?:-[a-z0-9]+)+-\d+$`)
+
 // column returns the CSV column declared in the YAML for the given
 // mapping key, or "" when the key is not declared. Required and
 // optional mappings are checked together; a mapping is valid wherever
@@ -332,7 +336,11 @@ func (s *GenericInventorySource) parseResourceRow(
 		if extra == nil {
 			extra = make(map[string]string)
 		}
-		extra[key] = cols.col(row, col)
+		value := cols.col(row, col)
+		if s.config.ID == "opensearch" && key == "region" {
+			value = normalizeOpenSearchRegion(resourceID, value)
+		}
+		extra[key] = value
 	}
 
 	// Service derivation: prefer the configured app tag; if none is
@@ -365,6 +373,20 @@ func (s *GenericInventorySource) parseResourceRow(
 	return resource, nil
 }
 
+func normalizeOpenSearchRegion(resourceID, region string) string {
+	if awsRegionPattern.MatchString(region) {
+		return region
+	}
+
+	resourceARN, err := arn.Parse(resourceID)
+	if err != nil || resourceARN.Service != "es" || !strings.HasPrefix(resourceARN.Resource, "domain/") ||
+		!awsRegionPattern.MatchString(resourceARN.Region) {
+		return region
+	}
+
+	return resourceARN.Region
+}
+
 // getReportIDFromMap reads the WIZ_REPORT_IDS JSON map and returns the report ID for the given resource
 func getReportIDFromMap(resourceID string) (string, error) {
 	// Read WIZ_REPORT_IDS environment variable
diff --git a/pkg/inventory/wiz/generic_test.go b/pkg/inventory/wiz/generic_test.go
index 081664e..e0096e4 100644
--- a/pkg/inventory/wiz/generic_test.go
+++ b/pkg/inventory/wiz/generic_test.go
@@ -803,6 +803,76 @@ func auroraTransforms() config.TransformsConfig {
 	}
 }
 
+func TestParseResourceRow_OpenSearchRegion(t *testing.T) {
+	cfg := config.ResourceConfig{
+		ID:            "opensearch",
+		Type:          "opensearch",
+		CloudProvider: "aws",
+		Inventory: config.InventoryConfig{
+			RequiredMappings: map[string]string{
+				"resource_id": "externalId",
+				"version":     "versionDetails.version",
+			},
+			FieldMappings: map[string]string{
+				"region": "region",
+			},
+		},
+	}
+
+	source := NewGenericInventorySource(&Client{}, &cfg, nil, nil)
+	cols := buildColumnIndex([]string{
+		"externalId",
+		"versionDetails.version",
+		"regionLocation",
+	})
+	tests := []struct {
+		name       string
+		resourceID string
+		region     string
+		want       string
+	}{
+		{
+			name:       "country location falls back to ARN region",
+			resourceID: "arn:aws:es:us-west-2:123456789012:domain/customer-search",
+			region:     "US",
+			want:       "us-west-2",
+		},
+		{
+			name:       "AWS region remains unchanged",
+			resourceID: "arn:aws:es:us-west-2:123456789012:domain/customer-search",
+			region:     "eu-west-1",
+			want:       "eu-west-1",
+		},
+		{
+			name:       "country location remains when resource ID is not an OpenSearch ARN",
+			resourceID: "customer-search",
+			region:     "US",
+			want:       "US",
+		},
+		{
+			name:       "country location remains when ARN is not an OpenSearch domain",
+			resourceID: "arn:aws:es:us-west-2:123456789012:package/example",
+			region:     "US",
+			want:       "US",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			row := []string{
+				tt.resourceID,
+				"OpenSearch_2.11",
+				tt.region,
+			}
+
+			resource, err := source.parseResourceRow(context.Background(), cols, row)
+
+			require.NoError(t, err)
+			assert.Equal(t, tt.want, resource.Extra["region"])
+		})
+	}
+}
+
 func TestParseResourceRow_Lambda(t *testing.T) {
 	cfg := config.ResourceConfig{
 		ID:            "lambda",