From f7c125694ad6b200bafd0d77002e4b3e2033d4ed Mon Sep 17 00:00:00 2001
From: Dave Arnold <user@example.com>
Date: Wed, 8 Jul 2026 16:35:21 -0400
Subject: [PATCH] fix: replace hardcoded target role ARN in ec2.py with env
 vars

- Add patch_target_role_name variable (default: r-inf-patch) for the
  cross-account role name that must exist in each target account
- Worker Lambda now passes aws_partition and patch_target_role_name as
  env vars alongside existing DB/queue vars
- ec2.py builds the target ARN dynamically from env vars instead of
  hardcoding partition (aws-us-gov) and role name (r-inf-patch)
---
 lambda-worker/src/ec2.py |  7 +++++--
 lambda.tf                | 12 +++++++-----
 variables.tf             |  6 ++++++
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/lambda-worker/src/ec2.py b/lambda-worker/src/ec2.py
index a30b397..239beeb 100644
--- a/lambda-worker/src/ec2.py
+++ b/lambda-worker/src/ec2.py
@@ -1,5 +1,6 @@
 import boto3
 import logging
+import os
 from botocore.config import Config
 from botocore.exceptions import ClientError
 
@@ -18,8 +19,10 @@ def assumedRoleSession(account, region):
         botoconf = Config(region_name=region, retries={'max_attempts': 5, 'mode': 'standard'})
         sts = boto3.client('sts', region_name=region, config=botoconf)
         
-        sts_role_arn = f"arn:aws-us-gov:iam::{account}:role/r-inf-patch"
-        sts_session_name = f"patch_session_{account}"
+        sts_role_arn = "arn:{}:iam::{}:role/{}".format(
+            os.environ["aws_partition"], account, os.environ["patch_target_role_name"]
+        )
+        sts_session_name = "patch_session_{}".format(account)
         
         try:
             # CHANGED: Ensure this variable name matches the one used below
diff --git a/lambda.tf b/lambda.tf
index d0b737b..78fef7f 100644
--- a/lambda.tf
+++ b/lambda.tf
@@ -139,11 +139,13 @@ resource "aws_lambda_function" "worker" {
 
   environment {
     variables = {
-      input_queue_url = aws_sqs_queue.dispatch_lambda.url
-      dbEndpoint      = aws_rds_cluster.patch4.endpoint
-      dbUsername      = local.db_user_name
-      dbDatabase      = local.db_name
-      dbRegion        = local.region
+      input_queue_url        = aws_sqs_queue.dispatch_lambda.url
+      aws_partition          = local.partition
+      patch_target_role_name = var.patch_target_role_name
+      dbEndpoint             = aws_rds_cluster.patch4.endpoint
+      dbUsername             = local.db_user_name
+      dbDatabase             = local.db_name
+      dbRegion               = local.region
     }
   }
 
diff --git a/variables.tf b/variables.tf
index 7005d5c..e240490 100644
--- a/variables.tf
+++ b/variables.tf
@@ -6,6 +6,12 @@ variable "name_prefix" {
   default     = "csvd-patch40"
 }
 
+variable "patch_target_role_name" {
+  description = "Name of the IAM role that exists in each target account and is assumed by the worker Lambda to perform EC2 operations (e.g. 'r-inf-patch')."
+  type        = string
+  default     = "r-inf-patch"
+}
+
 # ─── Additional Tags ──────────────────────────────────────────────────────────
 
 variable "tags" {