cloud-connect.sh

#!/bin/bash

set -euo pipefail

echo "Starting AWS integration setup"

: ${INTEGRATION_NAME:="AWS discovery"}
: ${CASTAI_API_URL:="https://api.cast.ai"}
: ${ROLE_NAME:="castai-discovery-role"}
: ${STACKSET_NAME:="castai-discovery-roles"}
: ${CASTAI_INTEGRATION_SCOPE:="ALL"}
: ${CASTAI_INTEGRATION_SETTINGS:="{}"}
: ${AWS_PERMISSIONS:=""}
: ${AWS_MANAGED_POLICIES:=""}
: ${AWS_CUR_S3_BUCKET_NAME:=""}
: ${AWS_CUR_S3_REGION:=""}
: ${AWS_CUR_S3_ACCOUNT_ID:=""}
: ${AWS_FORCE_ACCOUNT_SCOPE:="false"}
: ${AWS_ACCOUNT_IDS:=""}
: ${AWS_EKS_K8S_SYNC_ENABLED:="true"}
: ${AWS_EKS_CLUSTER_ARNS:=""}

if [ -z "${CASTAI_API_KEY}" ]; then
  echo "Error: Required environment variable CASTAI_API_KEY is not set."
  exit 1
fi

if [ -z "${CASTAI_ORGANIZATION_ID}" ]; then
  echo "Error: Required environment variable CASTAI_ORGANIZATION_ID is not set."
  exit 1
fi

if [ -z "${AWS_CAST_USER_ARN}" ]; then
    echo "Error: Required environment variable AWS_CAST_USER_ARN is not set."
    exit 1
fi

# Validate CUR S3 bucket environment variables - both must be provided or none
if ([ -n "${AWS_CUR_S3_BUCKET_NAME}" ] && [ -z "${AWS_CUR_S3_REGION}" ]) || ([ -z "${AWS_CUR_S3_BUCKET_NAME}" ] && [ -n "${AWS_CUR_S3_REGION}" ]); then
    echo "Error: Both AWS_CUR_S3_BUCKET_NAME and AWS_CUR_S3_REGION must be provided together or neither should be provided."
    exit 1
fi

# Build reusable S3 read-only statement scoped to CUR bucket if provided
S3_READ_STATEMENT=""
if [ -n "${AWS_CUR_S3_BUCKET_NAME}" ] && [ -n "${AWS_CUR_S3_REGION}" ]; then
    S3_READ_STATEMENT=$(cat <<EOF
{
    "Effect": "Allow",
    "Action": [
        "s3:Get*",
        "s3:List*"
    ],
    "Resource": [
        "arn:aws:s3:::${AWS_CUR_S3_BUCKET_NAME}",
        "arn:aws:s3:::${AWS_CUR_S3_BUCKET_NAME}/*"
    ]
}
EOF
)
fi

convert_to_json_array() {
    local input="$1"
    local default="$2"
    
    if [ -n "$input" ]; then
        echo "$input" | tr ',' '\n' | sed 's/^/"/;s/$/"/' | tr '\n' ',' | sed 's/,$//' | sed 's/^/[/;s/$/]/'
    else
        echo "$default"
    fi
}

PERMISSIONS_JSON=$(convert_to_json_array "${AWS_PERMISSIONS}" "")
# Convert comma-separated managed policy names to ARNs and then to JSON array
MANAGED_POLICIES_JSON="[]"
if [ -n "${AWS_MANAGED_POLICIES}" ]; then
    # Convert policy names to ARNs
    MANAGED_POLICY_ARNS=$(echo "$AWS_MANAGED_POLICIES" | tr ',' '\n' | sed 's/^/arn:aws:iam::aws:policy\//' | tr '\n' ',' | sed 's/,$//')
    # Convert ARNs to JSON array
    MANAGED_POLICIES_JSON=$(convert_to_json_array "$MANAGED_POLICY_ARNS" "[]")
fi

# Check if we're in a management account
if [ "$AWS_FORCE_ACCOUNT_SCOPE" = "true" ]; then
    echo "AWS_FORCE_ACCOUNT_SCOPE is set to true. Skipping management account detection and setting up account-scoped integration."
    ORG_SCOPE=false
else
    echo "Checking if current account is an AWS management account..."
    MANAGEMENT_ACCOUNT_ID=""
    if aws organizations describe-organization --no-cli-pager &>/dev/null; then
        MANAGEMENT_ACCOUNT_ID=$(aws organizations describe-organization --query 'Organization.MasterAccountId' --output text --no-cli-pager)
        CURRENT_ACCOUNT_ID=$(aws sts get-caller-identity --query Account --output text)
        
        if [ "$MANAGEMENT_ACCOUNT_ID" == "$CURRENT_ACCOUNT_ID" ]; then
            echo "Detected management account. Setting up org-scoped integration."
            ORG_SCOPE=true
        else
            echo "Not a management account. Setting up account-scoped integration."
            ORG_SCOPE=false
        fi
    else
        echo "No AWS Organizations access or not in an organization. Setting up account-scoped integration."
        ORG_SCOPE=false
    fi
fi

if [ "$ORG_SCOPE" = true ]; then
    echo "Setting up org-scoped integration with service-managed stackset..."

    DISCOVERY_POLICY_NAME="castai-discovery-readonly-policy"

    # Validate that we have either custom permissions or managed policies
    if [ -z "${AWS_PERMISSIONS}" ] && [ -z "${AWS_MANAGED_POLICIES}" ]; then
        echo "Error: Either AWS_PERMISSIONS or AWS_MANAGED_POLICIES environment variable is required."
        exit 1
    fi

    # Create IAM role in management account for org-wide operations
    echo "Creating management account role: $ROLE_NAME"
    
    MANAGEMENT_TRUST_POLICY=$(cat <<EOF
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": "${AWS_CAST_USER_ARN}"
            },
            "Action": "sts:AssumeRole",
            "Condition": {
                "StringEquals": {
                    "sts:ExternalId": "${CASTAI_ORGANIZATION_ID}"
                }
            }
        }
    ]
}
EOF
)
    
    if aws iam get-role --role-name "$ROLE_NAME" --no-cli-pager &>/dev/null; then
        echo "Management role '$ROLE_NAME' already exists. Updating trust policy."
        MANAGEMENT_ROLE_ARN=$(aws iam get-role --role-name "$ROLE_NAME" --query 'Role.Arn' --output text --no-cli-pager)
        aws iam update-assume-role-policy \
            --role-name "$ROLE_NAME" \
            --policy-document "${MANAGEMENT_TRUST_POLICY}" \
            --no-cli-pager
    else
        echo "Creating management role '$ROLE_NAME'."
        MANAGEMENT_ROLE_ARN=$(aws iam create-role \
            --role-name "$ROLE_NAME" \
            --assume-role-policy-document "${MANAGEMENT_TRUST_POLICY}" \
            --query 'Role.Arn' --output text --no-cli-pager)
    fi
    
    # Create management account policy for org operations
    MANAGEMENT_POLICY_NAME="castai-org-management-policy"
    MANAGEMENT_POLICY=$(cat <<EOF
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "organizations:ListAccounts",
                "organizations:DescribeOrganization"
            ],
            "Resource": "*"
        }
    ]
}
EOF
)
    
    MANAGEMENT_POLICY_ARN="arn:aws:iam::${CURRENT_ACCOUNT_ID}:policy/${MANAGEMENT_POLICY_NAME}"
    echo "Checking for management policy: $MANAGEMENT_POLICY_NAME"
    
    if ! aws iam get-policy --policy-arn "$MANAGEMENT_POLICY_ARN" --no-cli-pager &>/dev/null; then
        echo "Creating management policy '$MANAGEMENT_POLICY_NAME'."
        aws iam create-policy \
            --policy-name "$MANAGEMENT_POLICY_NAME" \
            --policy-document "$MANAGEMENT_POLICY" \
            --no-cli-pager > /dev/null
    else
        echo "Management policy '$MANAGEMENT_POLICY_NAME' already exists."
    fi
    
    # Attach management policy to management role
    if ! aws iam list-attached-role-policies --role-name "$ROLE_NAME" --no-cli-pager | grep -q "$MANAGEMENT_POLICY_NAME"; then
        echo "Attaching management policy to management role."
        aws iam attach-role-policy \
            --role-name "$ROLE_NAME" \
            --policy-arn "$MANAGEMENT_POLICY_ARN" \
            --no-cli-pager
    else
        echo "Management policy already attached to management role."
    fi

    # Determine if management account should have discovery permissions
    # Only attach discovery permissions if:
    # - AWS_ACCOUNT_IDS is empty (sync all accounts), OR
    # - AWS_ACCOUNT_IDS contains the management account ID (explicitly allowed)
    ATTACH_DISCOVERY_TO_MANAGEMENT=false
    if [ -z "${AWS_ACCOUNT_IDS}" ]; then
        echo "No account IDs specified - management account will have discovery permissions to sync all accounts."
        ATTACH_DISCOVERY_TO_MANAGEMENT=true
    else
        # Check if management account ID is in the list
        for account_id in $(echo "$AWS_ACCOUNT_IDS" | tr ',' ' '); do
            if [ "$account_id" = "$CURRENT_ACCOUNT_ID" ]; then
                echo "Management account ID found in AWS_ACCOUNT_IDS - management account will have discovery permissions."
                ATTACH_DISCOVERY_TO_MANAGEMENT=true
                break
            fi
        done
        if [ "$ATTACH_DISCOVERY_TO_MANAGEMENT" = false ]; then
            echo "Management account ID not in AWS_ACCOUNT_IDS - management account will NOT have discovery permissions (only org management permissions)."
        fi
    fi

    # Attach discovery managed policies to management role if provided and allowed
    if [ "$ATTACH_DISCOVERY_TO_MANAGEMENT" = true ] && [ -n "${AWS_MANAGED_POLICIES}" ]; then
        echo "Attaching discovery managed policies to management role: ${AWS_MANAGED_POLICIES}"
        echo "$AWS_MANAGED_POLICIES" | tr ',' '\n' | while read -r policy_name; do
            if [ -n "$policy_name" ]; then
                policy_arn="arn:aws:iam::aws:policy/$policy_name"
                echo "Checking if managed policy is attached: $policy_arn"
                if ! aws iam list-attached-role-policies --role-name "$ROLE_NAME" --no-cli-pager | grep -q "$(basename $policy_arn)"; then
                    echo "Attaching managed policy to management role: $policy_arn"
                    aws iam attach-role-policy \
                        --role-name "$ROLE_NAME" \
                        --policy-arn "$policy_arn" \
                        --no-cli-pager
                else
                    echo "Managed policy already attached: $policy_arn"
                fi
            fi
        done
    fi

    # Create and attach discovery custom policy to management role if permissions are provided and allowed
    if [ "$ATTACH_DISCOVERY_TO_MANAGEMENT" = true ] && [ -n "${AWS_PERMISSIONS}" ]; then
        DISCOVERY_CUSTOM_POLICY=$(cat <<EOF
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": $PERMISSIONS_JSON,
            "Resource": "*"
        }$(if [ -n "${S3_READ_STATEMENT}" ]; then echo ", ${S3_READ_STATEMENT}"; fi)
    ]
}
EOF
)
        DISCOVERY_CUSTOM_POLICY_ARN="arn:aws:iam::${CURRENT_ACCOUNT_ID}:policy/${DISCOVERY_POLICY_NAME}"
        echo "Checking for discovery custom policy: $DISCOVERY_POLICY_NAME"

        if ! aws iam get-policy --policy-arn "$DISCOVERY_CUSTOM_POLICY_ARN" --no-cli-pager &>/dev/null; then
            echo "Creating discovery custom policy '$DISCOVERY_POLICY_NAME'."
            aws iam create-policy \
                --policy-name "$DISCOVERY_POLICY_NAME" \
                --policy-document "$DISCOVERY_CUSTOM_POLICY" \
                --no-cli-pager > /dev/null
        else
            echo "Discovery custom policy '$DISCOVERY_POLICY_NAME' already exists. Updating with new version."

            # Delete old non-default versions to avoid hitting the 5-version limit
            OLD_VERSIONS=$(aws iam list-policy-versions --policy-arn "$DISCOVERY_CUSTOM_POLICY_ARN" --no-cli-pager --query 'Versions[?!IsDefaultVersion].VersionId' --output text)
            for version in $OLD_VERSIONS; do
                echo "Deleting old policy version: $version"
                aws iam delete-policy-version \
                    --policy-arn "$DISCOVERY_CUSTOM_POLICY_ARN" \
                    --version-id "$version" \
                    --no-cli-pager
            done

            # Create new policy version and set as default
            echo "Creating new policy version for '$DISCOVERY_POLICY_NAME'."
            aws iam create-policy-version \
                --policy-arn "$DISCOVERY_CUSTOM_POLICY_ARN" \
                --policy-document "$DISCOVERY_CUSTOM_POLICY" \
                --set-as-default \
                --no-cli-pager > /dev/null
        fi

        # Attach discovery custom policy to management role
        echo "Checking if discovery policy '$DISCOVERY_POLICY_NAME' is attached to management role."
        if ! aws iam list-attached-role-policies --role-name "$ROLE_NAME" --no-cli-pager | grep -q "$DISCOVERY_POLICY_NAME"; then
            echo "Attaching discovery policy '$DISCOVERY_POLICY_NAME' to management role."
            aws iam attach-role-policy \
                --role-name "$ROLE_NAME" \
                --policy-arn "$DISCOVERY_CUSTOM_POLICY_ARN" \
                --no-cli-pager
        else
            echo "Discovery policy '$DISCOVERY_POLICY_NAME' already attached to management role."
        fi
    fi

    # Generate Lambda resources for EKS access configuration if k8s sync is enabled
    EKS_LAMBDA_RESOURCES=""
    INCLUDE_EKS_LAMBDA=false
    if [ "$AWS_EKS_K8S_SYNC_ENABLED" = "true" ] && ([ "$CASTAI_INTEGRATION_SCOPE" = "ALL" ] || [ "$CASTAI_INTEGRATION_SCOPE" = "ALL_MINIMAL" ]); then
        echo "Including EKS access Lambda in StackSet template for automatic cluster access configuration..."
        INCLUDE_EKS_LAMBDA=true

        # Lambda code (Python) - handles CloudFormation Custom Resource to configure EKS access on stack deployment
        LAMBDA_CODE='import boto3
import json
import urllib.request
import os

def send_cfn_response(event, context, status, data=None):
    body = json.dumps({"Status": status, "Reason": "See CloudWatch", "PhysicalResourceId": context.log_stream_name, "StackId": event["StackId"], "RequestId": event["RequestId"], "LogicalResourceId": event["LogicalResourceId"], "Data": data or {}}).encode()
    req = urllib.request.Request(event["ResponseURL"], data=body, headers={"Content-Type": ""}, method="PUT")
    urllib.request.urlopen(req)

def get_allowed_arns():
    allowed = os.environ.get("ALLOWED_CLUSTER_ARNS", "")
    return set(allowed.split(",")) if allowed else None

def is_cluster_allowed(cluster_arn, allowed_arns):
    if allowed_arns is None:
        return True
    return cluster_arn in allowed_arns

def handler(event, context):
    print(f"Event: {json.dumps(event)}")
    if event.get("RequestType") == "Delete":
        send_cfn_response(event, context, "SUCCESS")
        return
    try:
        result = configure_all_clusters(os.environ["CAST_ROLE_ARN"])
        send_cfn_response(event, context, "SUCCESS", result)
    except Exception as e:
        print(f"Error: {e}")
        send_cfn_response(event, context, "FAILED", {"Error": str(e)})

def configure_all_clusters(role_arn):
    ec2 = boto3.client("ec2")
    sts = boto3.client("sts")
    account_id = sts.get_caller_identity()["Account"]
    allowed_arns = get_allowed_arns()
    configured, failed, skipped = [], [], []
    regions = [r["RegionName"] for r in ec2.describe_regions()["Regions"]]
    print(f"Scanning {len(regions)} regions for EKS clusters")
    for region in regions:
        try:
            eks = boto3.client("eks", region_name=region)
            for page in eks.get_paginator("list_clusters").paginate():
                for c in page["clusters"]:
                    cluster_arn = f"arn:aws:eks:{region}:{account_id}:cluster/{c}"
                    if not is_cluster_allowed(cluster_arn, allowed_arns):
                        print(f"Skipping cluster {c} - ARN {cluster_arn} not in allowed list")
                        skipped.append(cluster_arn)
                        continue
                    if configure_cluster(c, region, role_arn):
                        configured.append(cluster_arn)
                    else:
                        failed.append(cluster_arn)
        except Exception as e:
            print(f"Error scanning region {region}: {e}")
    return {"configured": configured, "failed": failed, "skipped": skipped}

def configure_cluster(name, region, role_arn):
    eks = boto3.client("eks", region_name=region)
    policy = "arn:aws:eks::aws:cluster-access-policy/AmazonEKSAdminViewPolicy"
    try:
        eks.describe_access_entry(clusterName=name, principalArn=role_arn)
        print(f"Access entry exists for {name}")
    except eks.exceptions.ResourceNotFoundException:
        try:
            eks.create_access_entry(clusterName=name, principalArn=role_arn, type="STANDARD")
            print(f"Created access entry for {name}")
        except Exception as e:
            print(f"Failed to create access entry for {name}: {e}")
            return False
    except Exception as e:
        print(f"Skipping {name}: {e}")
        return False
    try:
        resp = eks.list_associated_access_policies(clusterName=name, principalArn=role_arn)
        if policy not in [p["policyArn"] for p in resp.get("associatedAccessPolicies", [])]:
            eks.associate_access_policy(clusterName=name, principalArn=role_arn, policyArn=policy, accessScope={"type": "cluster"})
            print(f"Associated policy for {name}")
    except Exception as e:
        print(f"Failed to associate policy for {name}: {e}")
        return False
    return True
'

        # Escape the Lambda code for JSON embedding (escape backslashes, quotes, and newlines)
        LAMBDA_CODE_ESCAPED=$(echo "$LAMBDA_CODE" | python3 -c 'import sys,json; print(json.dumps(sys.stdin.read()))')

        # Build Lambda environment variables - include ALLOWED_CLUSTER_ARNS if specified
        LAMBDA_ENV_VARS='{"CAST_ROLE_ARN": {"Fn::GetAtt": ["CastDiscoveryRole", "Arn"]}'
        if [ -n "${AWS_EKS_CLUSTER_ARNS}" ]; then
            LAMBDA_ENV_VARS="${LAMBDA_ENV_VARS}, \"ALLOWED_CLUSTER_ARNS\": \"${AWS_EKS_CLUSTER_ARNS}\""
            echo "Lambda will only configure EKS access for specified cluster ARNs: ${AWS_EKS_CLUSTER_ARNS}"
        fi
        LAMBDA_ENV_VARS="${LAMBDA_ENV_VARS}}"

        EKS_LAMBDA_RESOURCES=',
        "EksAccessLambdaRole": {
            "Type": "AWS::IAM::Role",
            "Properties": {
                "RoleName": {"Fn::Sub": "${RoleName}-eks-lambda"},
                "AssumeRolePolicyDocument": {
                    "Version": "2012-10-17",
                    "Statement": [{"Effect": "Allow", "Principal": {"Service": "lambda.amazonaws.com"}, "Action": "sts:AssumeRole"}]
                },
                "ManagedPolicyArns": ["arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"],
                "Policies": [{
                    "PolicyName": "eks-access-management",
                    "PolicyDocument": {
                        "Version": "2012-10-17",
                        "Statement": [{"Effect": "Allow", "Action": ["ec2:DescribeRegions", "eks:ListClusters", "eks:DescribeCluster", "eks:CreateAccessEntry", "eks:DescribeAccessEntry", "eks:AssociateAccessPolicy", "eks:ListAssociatedAccessPolicies"], "Resource": "*"}]
                    }
                }]
            }
        },
        "EksAccessLambda": {
            "Type": "AWS::Lambda::Function",
            "Properties": {
                "FunctionName": {"Fn::Sub": "${RoleName}-eks-access"},
                "Runtime": "python3.11",
                "Handler": "index.handler",
                "Timeout": 300,
                "MemorySize": 256,
                "Role": {"Fn::GetAtt": ["EksAccessLambdaRole", "Arn"]},
                "Environment": {"Variables": '"$LAMBDA_ENV_VARS"'},
                "Code": {"ZipFile": '"$LAMBDA_CODE_ESCAPED"'}
            },
            "DependsOn": ["CastDiscoveryRole"]
        },
        "EksAccessTrigger": {
            "Type": "Custom::EksAccessConfiguration",
            "Properties": {"ServiceToken": {"Fn::GetAtt": ["EksAccessLambda", "Arn"]}},
            "DependsOn": ["EksAccessLambda"]
        }'
    fi

    # Create CloudFormation template for member account roles
    MEMBER_ROLE_TEMPLATE=$(cat <<EOF
{
    "AWSTemplateFormatVersion": "2010-09-09",
    "Description": "CAST AI discovery role for member accounts - ${CASTAI_INTEGRATION_SCOPE} scope",
    "Parameters": {
        "CastUserArn": {
            "Type": "String",
            "Description": "CAST AI user ARN"
        },
        "OrganizationId": {
            "Type": "String",
            "Description": "CAST AI organization ID"
        },
        "RoleName": {
            "Type": "String",
            "Description": "IAM role name for discovery"
        }
    },
    "Resources": {
        "CastDiscoveryRole": {
            "Type": "AWS::IAM::Role",
            "Properties": {
                "RoleName": {"Ref": "RoleName"},
                "AssumeRolePolicyDocument": {
                    "Version": "2012-10-17",
                    "Statement": [
                        {
                            "Effect": "Allow",
                            "Principal": {
                                "AWS": {"Ref": "CastUserArn"}
                            },
                            "Action": "sts:AssumeRole",
                            "Condition": {
                                "StringEquals": {
                                    "sts:ExternalId": {"Ref": "OrganizationId"}
                                }
                            }
                        }
                    ]
                },
                "ManagedPolicyArns": $MANAGED_POLICIES_JSON
            }
        }$(if [ -n "${AWS_PERMISSIONS}" ]; then echo ',
        "PermissionsPolicy": {
            "Type": "AWS::IAM::Policy",
            "Properties": {
                "PolicyName": "'${DISCOVERY_POLICY_NAME}'",
                "Roles": [{"Ref": "CastDiscoveryRole"}],
                "PolicyDocument": {
                    "Version": "2012-10-17",
                    "Statement": [
                        {
                            "Effect": "Allow",
                            "Action": '$PERMISSIONS_JSON',
                            "Resource": "*"
                        }'$(if [ -n "${S3_READ_STATEMENT}" ]; then echo ", ${S3_READ_STATEMENT}"; fi)'
                    ]
                }
            }
        }'; fi)${EKS_LAMBDA_RESOURCES}
    },
    "Outputs": {
        "RoleArn": {
            "Description": "ARN of the created role",
            "Value": {"Fn::GetAtt": ["CastDiscoveryRole", "Arn"]}
        }
    }
}
EOF
)
    
    # Create stackset for member account roles
    echo "Creating CloudFormation stackset: $STACKSET_NAME"
    
    # Check if stackset already exists
    if aws cloudformation describe-stack-set --stack-set-name "$STACKSET_NAME" --no-cli-pager &>/dev/null; then
        echo "Stackset '$STACKSET_NAME' already exists. Updating..."
        aws cloudformation update-stack-set \
            --stack-set-name "$STACKSET_NAME" \
            --template-body "$MEMBER_ROLE_TEMPLATE" \
            --parameters ParameterKey=CastUserArn,ParameterValue="$AWS_CAST_USER_ARN" \
                        ParameterKey=OrganizationId,ParameterValue="$CASTAI_ORGANIZATION_ID" \
                        ParameterKey=RoleName,ParameterValue="$ROLE_NAME" \
            --permission-model SERVICE_MANAGED \
            --auto-deployment Enabled=true,RetainStacksOnAccountRemoval=false \
            --capabilities CAPABILITY_NAMED_IAM \
            --no-cli-pager
    else
        echo "Creating new stackset '$STACKSET_NAME'..."
        aws cloudformation create-stack-set \
            --stack-set-name "$STACKSET_NAME" \
            --template-body "$MEMBER_ROLE_TEMPLATE" \
            --parameters ParameterKey=CastUserArn,ParameterValue="$AWS_CAST_USER_ARN" \
                        ParameterKey=OrganizationId,ParameterValue="$CASTAI_ORGANIZATION_ID" \
                        ParameterKey=RoleName,ParameterValue="$ROLE_NAME" \
            --permission-model SERVICE_MANAGED \
            --auto-deployment Enabled=true,RetainStacksOnAccountRemoval=false \
            --capabilities CAPABILITY_NAMED_IAM \
            --no-cli-pager
    fi
    
    # Deploy stack instances to member accounts
    echo "Deploying stack instances to member accounts..."

    # Get root OU IDs (always needed for SERVICE_MANAGED StackSets)
    ROOT_OU_IDS=$(aws organizations list-roots --query 'Roots[].Id' --output text --no-cli-pager)
    ROOT_OU_IDS_ARRAY=$(echo $ROOT_OU_IDS | tr '\t' ',' | sed 's/^/[/;s/$/]/')

    if [ -n "${AWS_ACCOUNT_IDS}" ]; then
        # Filter out management account from the list since SERVICE_MANAGED StackSets cannot deploy to management account
        FILTERED_ACCOUNT_IDS=""
        for account_id in $(echo "$AWS_ACCOUNT_IDS" | tr ',' ' '); do
            if [ "$account_id" != "$CURRENT_ACCOUNT_ID" ]; then
                if [ -z "$FILTERED_ACCOUNT_IDS" ]; then
                    FILTERED_ACCOUNT_IDS="$account_id"
                else
                    FILTERED_ACCOUNT_IDS="$FILTERED_ACCOUNT_IDS,$account_id"
                fi
            fi
        done

        if [ -n "$FILTERED_ACCOUNT_IDS" ]; then
            # Use INTERSECTION to deploy only to accounts that are in both the OU AND the account list
            ACCOUNT_IDS_ARRAY=$(echo "$FILTERED_ACCOUNT_IDS" | tr ',' '\n' | sed 's/^/"/;s/$/"/' | tr '\n' ',' | sed 's/,$//' | sed 's/^/[/;s/$/]/')
            DEPLOYMENT_TARGETS="OrganizationalUnitIds=$ROOT_OU_IDS_ARRAY,Accounts=$ACCOUNT_IDS_ARRAY,AccountFilterType=INTERSECTION"
            echo "Targeting specified member accounts in root OUs: $ACCOUNT_IDS_ARRAY"

            aws cloudformation create-stack-instances \
                --stack-set-name "$STACKSET_NAME" \
                --deployment-targets $DEPLOYMENT_TARGETS \
                --regions $(aws ec2 describe-regions --query 'Regions[?RegionName==`us-east-1`].RegionName' --output text --no-cli-pager) \
                --operation-preferences MaxConcurrentPercentage=100 \
                --no-cli-pager
        else
            echo "Only management account specified in AWS_ACCOUNT_IDS. Skipping StackSet deployment (management account role already created)."
        fi
    else
        # Deploy to all member accounts via root OUs (no account filtering)
        DEPLOYMENT_TARGETS="OrganizationalUnitIds=$ROOT_OU_IDS_ARRAY"
        echo "Targeting all accounts via root OUs: $ROOT_OU_IDS_ARRAY"

        aws cloudformation create-stack-instances \
            --stack-set-name "$STACKSET_NAME" \
            --deployment-targets $DEPLOYMENT_TARGETS \
            --regions $(aws ec2 describe-regions --query 'Regions[?RegionName==`us-east-1`].RegionName' --output text --no-cli-pager) \
            --operation-preferences MaxConcurrentPercentage=100 \
            --no-cli-pager
    fi
    
    # Use management role ARN for the integration
    ROLE_ARN="$MANAGEMENT_ROLE_ARN"
    
else
    echo "Setting up account-scoped integration..."

    TRUST_POLICY=$(cat <<EOF
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": "${AWS_CAST_USER_ARN}"
            },
            "Action": "sts:AssumeRole",
            "Condition": {
                "StringEquals": {
                    "sts:ExternalId": "${CASTAI_ORGANIZATION_ID}"
                }
            }
        }
    ]
}
EOF
)
    
    echo "Checking for IAM role: $ROLE_NAME"
    if aws iam get-role --role-name "$ROLE_NAME" --no-cli-pager &>/dev/null; then
        echo "IAM role '$ROLE_NAME' already exists. Updating trust policy."
        ROLE_ARN=$(aws iam get-role --role-name "$ROLE_NAME" --query 'Role.Arn' --output text --no-cli-pager)

        aws iam update-assume-role-policy \
            --role-name "$ROLE_NAME" \
            --policy-document "${TRUST_POLICY}" \
            --no-cli-pager
    else
        echo "IAM role '$ROLE_NAME' not found. Creating new role."
        ROLE_ARN=$(aws iam create-role \
            --role-name "$ROLE_NAME" \
            --assume-role-policy-document "${TRUST_POLICY}" \
            --query 'Role.Arn' --output text --no-cli-pager)
    fi
    echo "IAM role ARN: $ROLE_ARN"

    ATTACH_POLICY_ARN=""
    POLICY_NAME=""

    if [ -z "${AWS_PERMISSIONS}" ] && [ -z "${AWS_MANAGED_POLICIES}" ]; then
        echo "Error: Either AWS_PERMISSIONS or AWS_MANAGED_POLICIES environment variable is required."
        exit 1
    fi

    # Attach managed policies if provided
    if [ -n "${AWS_MANAGED_POLICIES}" ]; then
        echo "Attaching managed policies: ${AWS_MANAGED_POLICIES}"
        echo "$AWS_MANAGED_POLICIES" | tr ',' '\n' | while read -r policy_name; do
            if [ -n "$policy_name" ]; then
                policy_arn="arn:aws:iam::aws:policy/$policy_name"
                echo "Attaching managed policy: $policy_arn"
                aws iam attach-role-policy \
                    --role-name "$ROLE_NAME" \
                    --policy-arn "$policy_arn" \
                    --no-cli-pager
            fi
        done
    fi

    # Create custom policy if permissions are provided
    if [ -n "${AWS_PERMISSIONS}" ]; then
        POLICY_NAME="castai-$(echo "$CASTAI_INTEGRATION_SCOPE" | tr '[:upper:]' '[:lower:]' | tr '_' '-')-readonly-policy"
        CUSTOM_POLICY=$(cat <<EOF
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": $PERMISSIONS_JSON,
            "Resource": "*"
        }$(if [ -n "${S3_READ_STATEMENT}" ]; then echo ", ${S3_READ_STATEMENT}"; fi)
    ]
}
EOF
)
        ACCOUNT_ID=$(aws sts get-caller-identity --query Account --output text)
        CUSTOM_POLICY_ARN="arn:aws:iam::${ACCOUNT_ID}:policy/${POLICY_NAME}"
        echo "Checking for custom policy: $POLICY_NAME"

        if ! aws iam get-policy --policy-arn "$CUSTOM_POLICY_ARN" --no-cli-pager &>/dev/null; then
            echo "Creating custom policy '$POLICY_NAME'."
            aws iam create-policy \
                --policy-name "$POLICY_NAME" \
                --policy-document "$CUSTOM_POLICY" \
                --no-cli-pager > /dev/null
        else
            echo "Custom policy '$POLICY_NAME' already exists. Updating with new version."

            # Delete old non-default versions to avoid hitting the 5-version limit
            OLD_VERSIONS=$(aws iam list-policy-versions --policy-arn "$CUSTOM_POLICY_ARN" --no-cli-pager --query 'Versions[?!IsDefaultVersion].VersionId' --output text)
            for version in $OLD_VERSIONS; do
                echo "Deleting old policy version: $version"
                aws iam delete-policy-version \
                    --policy-arn "$CUSTOM_POLICY_ARN" \
                    --version-id "$version" \
                    --no-cli-pager
            done

            # Create new policy version and set as default
            echo "Creating new policy version for '$POLICY_NAME'."
            aws iam create-policy-version \
                --policy-arn "$CUSTOM_POLICY_ARN" \
                --policy-document "$CUSTOM_POLICY" \
                --set-as-default \
                --no-cli-pager > /dev/null
        fi
        ATTACH_POLICY_ARN="$CUSTOM_POLICY_ARN"
    fi

    # Attach custom policy if it was created
    if [ -n "${ATTACH_POLICY_ARN:-}" ]; then
        POLICY_BASENAME=$(basename "$ATTACH_POLICY_ARN")
        echo "Checking if policy '$POLICY_BASENAME' is attached to role '$ROLE_NAME'."
        if ! aws iam list-attached-role-policies --role-name "$ROLE_NAME" --no-cli-pager | grep -q "$POLICY_BASENAME"; then
            echo "Attaching policy '$POLICY_BASENAME' to role."
            aws iam attach-role-policy \
                --role-name "$ROLE_NAME" \
                --policy-arn "$ATTACH_POLICY_ARN" \
                --no-cli-pager
        else
            echo "Policy '$POLICY_BASENAME' already attached."
        fi
    fi
fi

# Setup EKS cluster access for k8s resource discovery (requires opt-in and ALL/ALL_MINIMAL scope)
K8S_OBJECTS_SYNC_ENABLED=false
if [ "$AWS_EKS_K8S_SYNC_ENABLED" = "true" ] && ([ "$CASTAI_INTEGRATION_SCOPE" = "ALL" ] || [ "$CASTAI_INTEGRATION_SCOPE" = "ALL_MINIMAL" ]); then
    K8S_OBJECTS_SYNC_ENABLED=true

    if [ "$ORG_SCOPE" = true ]; then
        # For org-scoped integrations, Lambda deployed via StackSet handles EKS access configuration
        echo "EKS cluster access for member accounts will be configured automatically via Lambda (deployed by StackSet)"
        echo "Note: Only existing EKS clusters will be configured. For new clusters, re-run the StackSet deployment."
    else
        # For account-scoped integrations, configure EKS access directly
        echo "Configuring EKS cluster access for k8s resource discovery..."

        ALL_REGIONS=$(aws ec2 describe-regions --query 'Regions[].RegionName' --output text --no-cli-pager 2>/dev/null || echo "")

        if [ -n "$ALL_REGIONS" ]; then
            TOTAL_CLUSTERS_CONFIGURED=0

            if [ -n "${AWS_EKS_CLUSTER_ARNS}" ]; then
                echo "EKS cluster discovery limited to specified cluster ARNs: ${AWS_EKS_CLUSTER_ARNS}"
                SPECIFIED_CLUSTER_ARNS=$(echo "$AWS_EKS_CLUSTER_ARNS" | tr ',' ' ')
            else
                echo "Discovering all EKS clusters in all regions"
                SPECIFIED_CLUSTER_ARNS=""
            fi

            for region in $ALL_REGIONS; do
                CLUSTER_NAMES=$(aws eks list-clusters --region "$region" --query 'clusters[]' --output text --no-cli-pager 2>/dev/null || echo "")

                if [ -n "$CLUSTER_NAMES" ]; then
                    echo "Found EKS clusters in $region: $(echo $CLUSTER_NAMES | tr '\t' ' ')"

                    ACCOUNT_ID_FOR_ARN=$(aws sts get-caller-identity --query Account --output text 2>/dev/null || echo "")

                    for cluster_name in $CLUSTER_NAMES; do
                        CLUSTER_ARN="arn:aws:eks:${region}:${ACCOUNT_ID_FOR_ARN}:cluster/${cluster_name}"

                        if [ -n "$SPECIFIED_CLUSTER_ARNS" ]; then
                            CLUSTER_MATCHED=false
                            for specified_arn in $SPECIFIED_CLUSTER_ARNS; do
                                if [ "$CLUSTER_ARN" = "$specified_arn" ]; then
                                    CLUSTER_MATCHED=true
                                    break
                                fi
                            done
                            if [ "$CLUSTER_MATCHED" = false ]; then
                                echo "  Skipping cluster $cluster_name (ARN $CLUSTER_ARN not in AWS_EKS_CLUSTER_ARNS)"
                                continue
                            fi
                        fi

                        echo "  Configuring access for cluster: $cluster_name"

                        if aws eks describe-access-entry \
                            --cluster-name "$cluster_name" \
                            --principal-arn "$ROLE_ARN" \
                            --region "$region" \
                            --no-cli-pager &>/dev/null; then
                            echo "    Access entry already exists for $cluster_name"
                        else
                            if CREATE_ACCESS_ENTRY_OUTPUT=$(aws eks create-access-entry \
                                --cluster-name "$cluster_name" \
                                --principal-arn "$ROLE_ARN" \
                                --type STANDARD \
                                --region "$region" \
                                --no-cli-pager 2>&1); then
                                echo "    Created access entry for $cluster_name"
                            else
                                echo "    Warning: Failed to create access entry for $cluster_name: $CREATE_ACCESS_ENTRY_OUTPUT"
                                continue
                            fi
                        fi

                        POLICY_ARN="arn:aws:eks::aws:cluster-access-policy/AmazonEKSAdminViewPolicy"

                        ASSOCIATED_POLICIES=$(aws eks list-associated-access-policies \
                            --cluster-name "$cluster_name" \
                            --principal-arn "$ROLE_ARN" \
                            --region "$region" \
                            --query 'associatedAccessPolicies[].policyArn' \
                            --output text \
                            --no-cli-pager 2>/dev/null || echo "")

                        if echo "$ASSOCIATED_POLICIES" | grep -q "$POLICY_ARN"; then
                            echo "    Policy already associated for $cluster_name"
                        else
                            if ASSOCIATE_POLICY_OUTPUT=$(aws eks associate-access-policy \
                                --cluster-name "$cluster_name" \
                                --principal-arn "$ROLE_ARN" \
                                --policy-arn "$POLICY_ARN" \
                                --access-scope type=cluster \
                                --region "$region" \
                                --no-cli-pager 2>&1); then
                                echo "    Associated read-only policy for $cluster_name"
                                TOTAL_CLUSTERS_CONFIGURED=$((TOTAL_CLUSTERS_CONFIGURED + 1))
                            else
                                echo "    Warning: Failed to associate policy for $cluster_name: $ASSOCIATE_POLICY_OUTPUT"
                            fi
                        fi
                    done
                fi
            done

            if [ $TOTAL_CLUSTERS_CONFIGURED -gt 0 ]; then
                echo "Successfully configured EKS access for $TOTAL_CLUSTERS_CONFIGURED cluster(s)"
            else
                echo "No EKS clusters were configured (clusters may not support access entries or require manual setup)"
            fi
        else
            echo "Warning: Could not list AWS regions. Skipping EKS cluster access configuration."
        fi
    fi
else
    echo "Skipping EKS cluster access configuration (requires AWS_EKS_K8S_SYNC_ENABLED=true and ALL/ALL_MINIMAL scope)"
fi

echo "Creating the cloud asset integration in CAST AI for scope: ${CASTAI_INTEGRATION_SCOPE}"

# Prepare metadata based on scope
METADATA_JSON="{\"crossRoleUserArn\": \"$AWS_CAST_USER_ARN\"}"
if [ "$ORG_SCOPE" = true ]; then
    METADATA_JSON="{\"crossRoleUserArn\": \"$AWS_CAST_USER_ARN\", \"organizationScope\": true}"
fi

# Add CUR S3 bucket information if provided
if [ -n "${AWS_CUR_S3_BUCKET_NAME}" ] && [ -n "${AWS_CUR_S3_REGION}" ]; then
    CUR_S3_BUCKET_JSON=", \"curS3Bucket\": {\"name\": \"$AWS_CUR_S3_BUCKET_NAME\", \"region\": \"$AWS_CUR_S3_REGION\""
    if [ -n "${AWS_CUR_S3_ACCOUNT_ID}" ]; then
        CUR_S3_BUCKET_JSON="$CUR_S3_BUCKET_JSON, \"accountId\": \"$AWS_CUR_S3_ACCOUNT_ID\""
    fi
    CUR_S3_BUCKET_JSON="$CUR_S3_BUCKET_JSON}"
    METADATA_JSON=$(echo "$METADATA_JSON" | sed "s/}$/$CUR_S3_BUCKET_JSON}/")
fi

# Add account IDs if provided (only for org-scoped integrations)
if [ "$ORG_SCOPE" = true ] && [ -n "${AWS_ACCOUNT_IDS}" ]; then
    ACCOUNT_IDS_JSON=$(echo "$AWS_ACCOUNT_IDS" | tr ',' '\n' | sed 's/^/"/;s/$/"/' | tr '\n' ',' | sed 's/,$//' | sed 's/^/[/;s/$/]/')
    ACCOUNT_IDS_METADATA_JSON=", \"accountIds\": $ACCOUNT_IDS_JSON"
    METADATA_JSON=$(echo "$METADATA_JSON" | sed "s/}$/$ACCOUNT_IDS_METADATA_JSON}/")
fi

# Add k8s objects sync flag
if [ "$K8S_OBJECTS_SYNC_ENABLED" = true ]; then
    K8S_SYNC_METADATA_JSON=", \"k8sObjectsSyncEnabled\": true"
    METADATA_JSON=$(echo "$METADATA_JSON" | sed "s/}$/$K8S_SYNC_METADATA_JSON}/")
fi

# Add EKS cluster ARNs if provided
if [ -n "${AWS_EKS_CLUSTER_ARNS}" ]; then
    EKS_CLUSTER_ARNS_JSON=$(echo "$AWS_EKS_CLUSTER_ARNS" | tr ',' '\n' | sed 's/^/"/;s/$/"/' | tr '\n' ',' | sed 's/,$//' | sed 's/^/[/;s/$/]/')
    # Use | as delimiter since ARNs contain /
    METADATA_JSON="${METADATA_JSON%\}}, \"eksClusterArns\": $EKS_CLUSTER_ARNS_JSON}"
fi

RESPONSE=$(curl -s -w "\n%{http_code}" -X 'POST' \
  "$CASTAI_API_URL/inventory/v1beta/organizations/$CASTAI_ORGANIZATION_ID/cloud-asset-integrations" \
  -H 'accept: application/json' \
  -H "X-API-Key: $CASTAI_API_KEY" \
  -H 'Content-Type: application/json' \
  -d "{
  \"enabled\": true,
  \"name\": \"$INTEGRATION_NAME\",
  \"provider\": \"AWS\",
  \"scope\": \"$CASTAI_INTEGRATION_SCOPE\",
  \"aws_credentials\": {
    \"assume_role_arn\": \"$ROLE_ARN\"
  },
  \"settings\": $CASTAI_INTEGRATION_SETTINGS,
  \"metadata\": $METADATA_JSON
}")

STATUS_CODE=$(echo "$RESPONSE" | tail -n1)
RESPONSE_BODY=$(echo "$RESPONSE" | sed '$d')

if [ "$STATUS_CODE" -ne 200 ]; then
    echo "Error: Failed to create integration. Status code: $STATUS_CODE"
    echo "Response: $RESPONSE_BODY"
    exit 1
fi

echo "AWS integration setup is complete"
if [ "$ORG_SCOPE" = true ]; then
    echo "Organization-scoped integration created with service-managed stackset deployment."
    if [ "$INCLUDE_EKS_LAMBDA" = true ]; then
        echo "EKS access Lambda will be deployed to member accounts for automatic cluster access configuration."
    fi
else
    echo "Account-scoped integration created."
fi