Merge pull request #1 from terraform-modules/feature/migrate-csvd-module

feat: migrate CSVD patch4 module code from CSVD/terraform-csvd-patch4

data.tf

@@ -1,7 +1,49 @@
+data "aws_partition" "current" {}
 data "aws_caller_identity" "current" {}
+data "aws_region" "current" {}
+data "aws_organizations_organization" "org" {}
 
-data "aws_arn" "current" {
-  arn = data.aws_caller_identity.current.arn
+data "aws_db_subnet_group" "patch_db" {
+  name = var.db_subnet_group_name
 }
 
-data "aws_region" "current" {}
+data "aws_vpcs" "vpcs" {
+  filter {
+    name   = "tag:Name"
+    values = [var.vpc_name]
+  }
+}
+
+data "aws_vpc" "vpcs" {
+  for_each = toset(data.aws_vpcs.vpcs.ids)
+  id       = each.value
+}
+
+data "aws_security_group" "patch_db_sg" {
+  name = var.db_security_group_name
+  filter {
+    name   = "vpc-id"
+    values = [one(data.aws_vpcs.vpcs.ids)]
+  }
+}
+
+data "aws_security_group" "patch_lambda_sg" {
+  name = var.lambda_security_group_name
+  filter {
+    name   = "vpc-id"
+    values = [one(data.aws_vpcs.vpcs.ids)]
+  }
+}
+
+# Lambda functions must deploy into *-apps-* subnets
+data "aws_subnets" "apps" {
+  for_each = data.aws_vpc.vpcs
+  filter {
+    name   = "tag:Name"
+    values = ["*-apps-*"]
+  }
+  filter {
+    name   = "vpc-id"
+    values = [each.value.id]
+  }
+}

kms.tf

@@ -0,0 +1,6 @@
+module "patch4_key" {
+  source   = "git@github.e.it.census.gov:terraform-modules/aws-s3.git//kms_key?ref=tf-upgrade"
+  key_name = var.name_prefix
+
+  tags = merge(local.tags_kms, { Name = local.kms_alias_name })
+}

lambda-dispatcher/src/lambda_function.py

@@ -0,0 +1,74 @@
+import json
+import logging
+import os
+import boto3
+from botocore.exceptions import ClientError
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+# Load the two queue URLs into a hash (dictionary)
+QUEUE_MAP = {
+    1: os.environ.get('TARGET_SQS_QUEUE_LAMBDA_URL'),           # Lambda Workers
+    2: os.environ.get('TARGET_SQS_QUEUE_REMOTEEXECUTION_URL'), # Remote Execution
+}
+
+# Validate that all queues are configured
+for key, url in QUEUE_MAP.items():
+    if not url:
+        raise ValueError(f"Environment variable for dispatch_type {key} is not set")
+
+sqs_client = boto3.client('sqs')
+
+def lambda_handler(event, context):
+    logger.debug(f"Passed event: {event}")
+    """
+    Dispatches task to the correct SQS queue based on dispatch_type.
+    Expects a flattened payload from Aurora.
+    """
+    try:
+        # Aurora sends the JSON directly as the event dict
+        payload = event if isinstance(event, dict) else json.loads(event)
+
+        uid = payload.get('uid')
+        dispatch_type = payload.get('dispatch_type')
+
+        if not uid:
+            raise ValueError("Missing required field: uid")
+        if dispatch_type not in QUEUE_MAP:
+            raise ValueError(f"Invalid dispatch_type: {dispatch_type}. Must be 1, 2, or 3.")
+
+        queue_url = QUEUE_MAP[dispatch_type]
+
+        logger.info(f"Dispatching task - uid={uid} | dispatch_type={dispatch_type} | queue={queue_url}")
+
+        # ADJUSTMENT: Since the procedure is flattened, 'payload' IS the work data.
+        # We send the entire payload to SQS so the Worker sees all columns.
+        message_body = payload 
+
+        # Send to the correct SQS queue
+        response = sqs_client.send_message(
+            QueueUrl=queue_url,
+            MessageBody=json.dumps(message_body)
+        )
+
+        message_id = response.get('MessageId')
+        if not message_id:
+            raise ValueError("SQS returned no MessageId")
+
+        logger.info(f"Successfully enqueued task uid={uid} to queue_url={queue_url} → MessageId={message_id}")
+
+        return {
+            "success": True,
+            "messageId": message_id
+        }
+
+    except ClientError as e:
+        error_msg = f"SQS ClientError: {e.response['Error']['Message']}"
+        logger.error(error_msg, exc_info=True)
+        return {"success": False, "error": error_msg}
+
+    except Exception as e:
+        error_msg = f"Unexpected error dispatching task: {str(e)}"
+        logger.exception(error_msg)
+        return {"success": False, "error": error_msg}

lambda-layer/requirements.txt

@@ -0,0 +1,5 @@
+aws_advanced_python_wrapper==2.1.0
+psycopg==3.2.13
+psycopg-binary==3.2.13
+psycopg-pool==3.3.0
+typing_extensions

lambda-step-result/src/lambda_function.py

@@ -0,0 +1,102 @@
+import boto3
+import logging
+import json
+import psycopg
+import os
+from psycopg import OperationalError, InterfaceError
+import time
+
+# Enhanced Logging Setup
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG) 
+
+# Persistent connection variable
+db_connection = None
+
+def get_iam_token():
+    # RDS client to generate the auth token
+    client = boto3.client('rds', region_name=os.environ.get('AWS_REGION', os.getenv("dbRegion")))
+
+    return client.generate_db_auth_token(
+        DBHostname=os.getenv("dbEndpoint"),
+        Port=5432,
+        DBUsername=os.getenv("dbUsername"),
+        Region=os.getenv("dbRegion")
+    )
+
+def get_connection():
+    global db_connection
+
+    # Check if existing connection is alive
+    if db_connection is not None:
+        if db_connection.closed or db_connection.broken:
+            db_connection = None
+        else:
+            try:
+                db_connection.execute("SELECT 1")
+                return db_connection
+            except (OperationalError, Exception):
+                db_connection = None
+
+    # Retry Logic Configuration
+    max_retries = 3
+    retry_delay = 0.5  # Initial delay in seconds
+
+    for attempt in range(max_retries):
+        try:
+            print(f"Connection attempt {attempt + 1}...")
+
+            db_connection = psycopg.connect(
+                host=os.getenv("dbEndpoint"),
+                dbname=os.getenv("dbDatabase"),
+                user=os.getenv("dbUsername"),
+                password=get_iam_token(),
+                sslmode='require',
+                connect_timeout=5  # How long to wait for the TCP handshake
+            )
+            db_connection.autocommit = True
+            print("Connection established successfully.")
+            return db_connection
+
+        except OperationalError as e:
+            print(f"Attempt {attempt + 1} failed: {e}")
+            if attempt < max_retries - 1:
+                time.sleep(retry_delay)
+                retry_delay *= 2  # Exponential backoff (0.5s, 1s, 2s)
+            else:
+                raise e  # Exhausted retries
+def executeSql(sql, sqldata = None):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute(sql, sqldata)
+    conn.commit()
+    cursor.close()
+    conn.close()
+
+def executeSqlWithReturn(sql, sqldata = None):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute(sql, sqldata)
+    columns = [i[0] for i in cursor.description]
+    result = [dict(zip(columns, row)) for row in cursor.fetchall()]
+    conn.commit()
+    cursor.close()
+    conn.close()
+    return result
+
+def lambda_handler(event, context):
+    logger.debug(json.dumps(event))
+    body = json.loads(event['Records'][0]['body'])
+    #logger.debug(f"Processing body: {body}")
+
+    try:
+        loggit = executeSql("CALL patch40.sp_step_update(%s::bigint, %s::integer)",
+        (int(body['uid']), int(body['status_code'])))
+    except Exception as e:
+        raise e 
+    return {
+        'statusCode': 200,
+        'Payload': {
+            'status': 'SUCCESS'
+        }
+    }

lambda-worker/src/ec2.py

@@ -0,0 +1,123 @@
+import boto3
+import logging
+import os
+from botocore.config import Config
+from botocore.exceptions import ClientError
+
+logger = logging.getLogger(__name__)
+
+# Global session cache for warm-start reuse
+_session_cache = {}
+
+class ec2:
+    @staticmethod
+    def assumedRoleSession(account, region):
+        cache_key = f"{account}-{region}"
+        if cache_key in _session_cache:
+            return _session_cache[cache_key]
+
+        botoconf = Config(region_name=region, retries={'max_attempts': 5, 'mode': 'standard'})
+        sts = boto3.client('sts', region_name=region, config=botoconf)
+
+        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
+            sts_response = sts.assume_role(
+                RoleArn=sts_role_arn, 
+                RoleSessionName=sts_session_name
+            )
+
+            new_session = boto3.Session(
+                region_name=region,
+                aws_access_key_id=sts_response['Credentials']['AccessKeyId'],
+                aws_secret_access_key=sts_response['Credentials']['SecretAccessKey'],
+                aws_session_token=sts_response['Credentials']['SessionToken']
+            )
+
+            _session_cache[cache_key] = new_session
+            return new_session
+        except Exception as e:
+            logger.error(f"Error assuming role for {account}: {e}")
+            raise
+
+    @staticmethod
+    def getStopProtection(event):
+        try:
+            session = ec2.assumedRoleSession(account=event['accountid'], region=event['region'])
+            client = session.client('ec2')
+            response = client.describe_instance_attribute(
+                Attribute='disableApiStop', 
+                InstanceId=event['instanceid']
+            )
+            state = response['DisableApiStop']['Value']
+            return 2925 if state else 2930 # 2925=Enabled, 2930=Disabled
+        except ClientError as e:
+            return 4920 if "does not exist" in str(e) else 4000
+        except Exception:
+            return 4000
+
+    @staticmethod
+    def getRunState(event):
+        try:
+            session = ec2.assumedRoleSession(account=event['accountid'], region=event['region'])
+            client = session.client('ec2')
+            response = client.describe_instances(InstanceIds=[event['instanceid']])
+            state = response['Reservations'][0]['Instances'][0]['State']['Name']
+            return 2905 if state == "running" else 2910
+        except ClientError as e:
+            return 4920 if "does not exist" in str(e) else 4000
+        except Exception:
+            return 4000
+
+    @staticmethod
+    def startInstance(event):
+        try:
+            session = ec2.assumedRoleSession(account=event['accountid'], region=event['region'])
+            client = session.client('ec2')
+            client.start_instances(InstanceIds=[event['instanceid']])
+            return 2000 # Success
+        except ClientError as e:
+            return 4920 if "does not exist" in str(e) else 4000
+        except Exception:
+            return 4000
+
+    @staticmethod
+    def stopInstance(event):
+        try:
+            session = ec2.assumedRoleSession(account=event['accountid'], region=event['region'])
+            client = session.client('ec2')
+            client.stop_instances(InstanceIds=[event['instanceid']])
+            return 2000 # Success
+        except ClientError as e:
+            return 4920 if "does not exist" in str(e) else 4000
+        except Exception as e:
+            logger.exception("Exception: %s" % (str(e)))
+            return 4000
+
+    @staticmethod
+    def setStopProtection(event):
+        try:
+            session = ec2.assumedRoleSession(account=event['accountid'], region=event['region'])
+            client = session.client('ec2')
+            client.modify_instance_attribute(DisableApiStop={'Value': True}, InstanceId=event['instanceid'])
+            return 2000 # Success
+        except ClientError as e:
+            return 4920 if "does not exist" in str(e) else 4000
+        except Exception:
+            return 4000
+
+    @staticmethod
+    def removeStopProtection(event):
+        try:
+            session = ec2.assumedRoleSession(account=event['accountid'], region=event['region'])
+            client = session.client('ec2')
+            client.modify_instance_attribute(DisableApiStop={'Value': False}, InstanceId=event['instanceid'])
+            return 2000 # Success
+        except ClientError as e:
+            return 4920 if "does not exist" in str(e) else 4000
+        except Exception:
+            return 4000