buildspec.yml

version: 0.2
env:
  variables:
    # Opt out of AWS SAM CLI telemetry collection for all phases.
    SAM_CLI_TELEMETRY: "0"
phases:
  install:
    runtime-versions:
      python: 3.12
    commands:
      - echo "Installing dependencies..."
      - |
        if ! pip install -q --upgrade pip; then
          echo "ERROR: Failed to upgrade pip"
          exit 1
        fi
      - |
        if ! pip install -q --upgrade aws-sam-cli; then
          echo "ERROR: Failed to install aws-sam-cli"
          exit 1
        fi
      - |
        # Verify jq is available (used for parsing JSON credentials)
        if ! command -v jq &> /dev/null; then
          echo "ERROR: jq is not installed"
          exit 1
        fi
  build:
    commands:
      - echo "Starting AIML Security Assessment"
      - |
        # Keep deployment, artifact checks, and consolidated report scope aligned.
        export ENABLE_BEDROCK="${ENABLE_BEDROCK:-true}"
        export ENABLE_SAGEMAKER="${ENABLE_SAGEMAKER:-true}"
        export ENABLE_AGENTCORE="${ENABLE_AGENTCORE:-true}"
        export ENABLE_AGENT_REGISTRY="${ENABLE_AGENT_REGISTRY:-true}"
        for service_flag in ENABLE_BEDROCK ENABLE_SAGEMAKER ENABLE_AGENTCORE ENABLE_AGENT_REGISTRY; do
          case "${!service_flag}" in
            true|false) ;;
            *) echo "ERROR: $service_flag must be true or false"; exit 1 ;;
          esac
        done
      # EnableResponsibleAIGRCAssessment toggle (CodeBuild env var; default false).
      # Threaded into the Step Functions execution input as enableResponsibleAIGRC;
      # the ASL Responsible AI GRC branch runs only when this is "true".
      # Lower-cased here so the ASL's case-sensitive StringEquals "true" gate
      # matches a "True"/"TRUE" input.
      - export ENABLE_RESPONSIBLE_AI_GRC="$(echo "${ENABLE_RESPONSIBLE_AI_GRC:-false}" | tr '[:upper:]' '[:lower:]')"
      - |
        # EnableFinServAssessment legacy alias.
        #
        # EnableResponsibleAIGRCAssessment/ENABLE_RESPONSIBLE_AI_GRC always
        # materializes a value in CloudFormation (default "false"), so there is
        # no CFN-level way to distinguish "customer left this at its default"
        # from "customer explicitly chose false". EnableFinServAssessment
        # therefore uses the literal sentinel "__UNSET__" as its default so this
        # script CAN tell "alias not touched" apart from "alias explicitly set".
        #
        # Precedence:
        #   alias unset                                  -> primary value wins
        #   alias set, primary at its "false" default     -> alias value wins
        #   alias set, primary explicitly "true", agree    -> shared value
        #   alias set, primary explicitly "true", disagree -> explicit build failure
        ENABLE_FINSERV_RAW="${ENABLE_FINSERV:-__UNSET__}"
        if [[ "$ENABLE_FINSERV_RAW" == "__UNSET__" ]]; then
          echo "ENABLE_FINSERV not set; using ENABLE_RESPONSIBLE_AI_GRC ($ENABLE_RESPONSIBLE_AI_GRC)"
        else
          ENABLE_FINSERV_NORMALIZED="$(echo "$ENABLE_FINSERV_RAW" | tr '[:upper:]' '[:lower:]')"
          if [[ "$ENABLE_RESPONSIBLE_AI_GRC" == "true" && "$ENABLE_FINSERV_NORMALIZED" != "true" ]]; then
            echo "ERROR: EnableResponsibleAIGRCAssessment is explicitly \"true\" but EnableFinServAssessment is explicitly \"$ENABLE_FINSERV_NORMALIZED\"."
            echo "These are aliases for the same toggle and must agree. Set them to the same value, or leave EnableFinServAssessment unset."
            exit 1
          fi
          export ENABLE_RESPONSIBLE_AI_GRC="$ENABLE_FINSERV_NORMALIZED"
          echo "ENABLE_FINSERV set to $ENABLE_RESPONSIBLE_AI_GRC; overriding the effective Responsible AI GRC toggle"
        fi
        unset ENABLE_FINSERV_RAW ENABLE_FINSERV_NORMALIZED
      - echo "Responsible AI GRC assessment enabled is $ENABLE_RESPONSIBLE_AI_GRC"
      # EnableOWASPAssessment toggle (CodeBuild env var; default false). Threaded
      # into the Step Functions execution input as enableOWASP; the ASL OWASP
      # branch runs only when this is "true". Lower-cased here so the ASL's
      # case-sensitive StringEquals "true" gate matches a "True"/"TRUE" input.
      - export ENABLE_OWASP="$(echo "${ENABLE_OWASP:-false}" | tr '[:upper:]' '[:lower:]')"
      - echo "OWASP assessment enabled is $ENABLE_OWASP"
      - |
        # Validate required environment variables
        missing_vars=()
        [[ -z "$AWS_DEFAULT_REGION" ]] && missing_vars+=("AWS_DEFAULT_REGION")

        if [[ $MULTI_ACCOUNT_SCAN == 'true' ]]; then
          [[ -z "$BUCKET_REPORT" ]] && missing_vars+=("BUCKET_REPORT")
          [[ -z "$AWS_PARTITION" ]] && missing_vars+=("AWS_PARTITION")
          [[ -z "$AWS_ACCOUNT_ID" ]] && missing_vars+=("AWS_ACCOUNT_ID")
          [[ -z "$MEMBER_ROLE" ]] && missing_vars+=("MEMBER_ROLE")
        fi

        if [[ ${#missing_vars[@]} -gt 0 ]]; then
          echo "ERROR: Missing required environment variables: ${missing_vars[*]}"
          exit 1
        fi
      - |
        normalize_delimited_list() {
          local value="$1"
          value=$(printf '%s' "$value" | sed -E 's/[[:space:]]*,[[:space:]]*/,/g; s/^[[:space:]]+//; s/[[:space:]]+$//; s/[[:space:]]+/,/g; s/,+/,/g; s/^,//; s/,$//')
          printf '%s' "$value"
        }

        # Accept comma- or whitespace-delimited input from CloudFormation, then
        # normalize to comma-delimited values for SAM and Lambda configuration.
        TARGET_REGIONS=$(normalize_delimited_list "${TARGET_REGIONS:-}")
        export TARGET_REGIONS

        MULTI_ACCOUNT_LIST_OVERRIDE=$(normalize_delimited_list "${MULTI_ACCOUNT_LIST_OVERRIDE:-}")
        export MULTI_ACCOUNT_LIST_OVERRIDE

        # Validate the normalized value matches the CloudFormation AllowedPattern.
        region_pattern='^[a-z]{2}(-gov)?-[a-z]+-[0-9](,[a-z]{2}(-gov)?-[a-z]+-[0-9])*$'
        if [[ -n "$TARGET_REGIONS" && ! "$TARGET_REGIONS" =~ $region_pattern ]]; then
          echo "ERROR: TARGET_REGIONS contains invalid characters: '$TARGET_REGIONS'"
          echo "Expected a comma- or space-separated list of region names (e.g., us-east-1,us-west-2), or empty."
          exit 1
        fi
        echo "Using TARGET_REGIONS: '${TARGET_REGIONS:-<deployment region only>}'"

        account_list_pattern='^[0-9]{12}(,[0-9]{12})*$'
        if [[ -n "$MULTI_ACCOUNT_LIST_OVERRIDE" && ! "$MULTI_ACCOUNT_LIST_OVERRIDE" =~ $account_list_pattern ]]; then
          echo "ERROR: MULTI_ACCOUNT_LIST_OVERRIDE contains invalid account IDs: '$MULTI_ACCOUNT_LIST_OVERRIDE'"
          echo "Expected a comma- or space-separated list of 12-digit AWS account IDs, or empty."
          exit 1
        fi

        # SAM's legacy CLI parser drops an unquoted empty ParameterValue and
        # then reuses the inner stack's previous value. Preserve literal quotes
        # around clearable values so empty input is sent as ParameterValue="".
        SAM_TARGET_REGIONS_PARAMETER="ParameterKey=TargetRegions,ParameterValue=\"${TARGET_REGIONS}\""
        SAM_REQUIRE_BEDROCK_ZERO_DATA_RETENTION_PARAMETER="ParameterKey=RequireBedrockZeroDataRetention,ParameterValue=${REQUIRE_BEDROCK_ZERO_DATA_RETENTION:-false}"
        SAM_REQUIRE_MARKETPLACE_ENDPOINT_CMK_PARAMETER="ParameterKey=RequireMarketplaceEndpointCMK,ParameterValue=${REQUIRE_MARKETPLACE_ENDPOINT_CMK:-true}"
        SAM_REQUIRE_AGENTCORE_ONLINE_EVALUATION_PARAMETER="ParameterKey=RequireAgentCoreOnlineEvaluation,ParameterValue=${REQUIRE_AGENTCORE_ONLINE_EVALUATION:-false}"
        SAM_REQUIRE_AGENT_REGISTRY_CMK_PARAMETER="ParameterKey=RequireAgentRegistryCMK,ParameterValue=${REQUIRE_AGENT_REGISTRY_CMK:-false}"
        SAM_AGENTCORE_TOKEN_VAULT_ID_PARAMETER="ParameterKey=AgentCoreTokenVaultId,ParameterValue=${AGENTCORE_TOKEN_VAULT_ID:-default}"
        SAM_ENABLE_AGENTCORE_ARTIFACT_OBJECT_READS_PARAMETER="ParameterKey=EnableAgentCoreArtifactObjectReads,ParameterValue=${ENABLE_AGENTCORE_ARTIFACT_OBJECT_READS:-false}"
        SAM_APPROVED_EXTERNAL_ACCOUNT_IDS_PARAMETER="ParameterKey=ApprovedExternalAccountIds,ParameterValue=\"${APPROVED_EXTERNAL_ACCOUNT_IDS:-}\""
        SAM_APPROVED_ORGANIZATION_IDS_PARAMETER="ParameterKey=ApprovedOrganizationIds,ParameterValue=\"${APPROVED_ORGANIZATION_IDS:-}\""
        SAM_ENABLE_SAGEMAKER_ARTIFACT_OBJECT_READS_PARAMETER="ParameterKey=EnableSageMakerArtifactObjectReads,ParameterValue=${ENABLE_SAGEMAKER_ARTIFACT_OBJECT_READS:-false}"
        SAM_ENABLE_BEDROCK_PARAMETER="ParameterKey=EnableBedrockAssessment,ParameterValue=${ENABLE_BEDROCK}"
        export SAM_ENABLE_BEDROCK_PARAMETER
        SAM_ENABLE_SAGEMAKER_PARAMETER="ParameterKey=EnableSageMakerAssessment,ParameterValue=${ENABLE_SAGEMAKER}"
        export SAM_ENABLE_SAGEMAKER_PARAMETER
        SAM_ENABLE_AGENTCORE_PARAMETER="ParameterKey=EnableAgentCoreAssessment,ParameterValue=${ENABLE_AGENTCORE}"
        export SAM_ENABLE_AGENTCORE_PARAMETER
        SAM_ENABLE_AGENT_REGISTRY_PARAMETER="ParameterKey=EnableAgentRegistryAssessment,ParameterValue=${ENABLE_AGENT_REGISTRY}"
        export SAM_ENABLE_AGENT_REGISTRY_PARAMETER
        export SAM_TARGET_REGIONS_PARAMETER
        export SAM_REQUIRE_BEDROCK_ZERO_DATA_RETENTION_PARAMETER
        export SAM_REQUIRE_MARKETPLACE_ENDPOINT_CMK_PARAMETER
        export SAM_REQUIRE_AGENTCORE_ONLINE_EVALUATION_PARAMETER
        export SAM_REQUIRE_AGENT_REGISTRY_CMK_PARAMETER
        export SAM_AGENTCORE_TOKEN_VAULT_ID_PARAMETER
        export SAM_ENABLE_AGENTCORE_ARTIFACT_OBJECT_READS_PARAMETER
        export SAM_APPROVED_EXTERNAL_ACCOUNT_IDS_PARAMETER
        export SAM_APPROVED_ORGANIZATION_IDS_PARAMETER
        export SAM_ENABLE_SAGEMAKER_ARTIFACT_OBJECT_READS_PARAMETER
      - echo "Multi-account scan is $MULTI_ACCOUNT_SCAN"
      - cd aiml-security-assessment
      - ls -la
      - |
        if [[ $MULTI_ACCOUNT_SCAN = 'true' ]]; then
          echo "Building multi-account template"
          if ! sam build --template template-multi-account.yaml; then
            echo "ERROR: SAM build failed for multi-account template"
            exit 1
          fi
        else
          echo "Building single-account template"
          if ! sam build --template template.yaml; then
            echo "ERROR: SAM build failed for single-account template"
            exit 1
          fi
        fi
      - echo "Build completed, checking build directory:"
      - ls -la .aws-sam/build/
      - |
        recover_failed_stack() {
          local stack_name="$1"
          local stack_status
          stack_status=$(aws cloudformation describe-stacks --stack-name "$stack_name" --query 'Stacks[0].StackStatus' --output text 2>/dev/null)
          if [[ "$stack_status" == "ROLLBACK_COMPLETE" || "$stack_status" == "DELETE_FAILED" ]]; then
            echo "Cleaning up failed deployment stack $stack_name (status: $stack_status)"
            if ! aws cloudformation delete-stack --stack-name "$stack_name"; then
              echo "ERROR: Failed to delete deployment stack $stack_name"
              return 1
            fi
            if ! aws cloudformation wait stack-delete-complete --stack-name "$stack_name"; then
              echo "ERROR: Timed out waiting for deployment stack $stack_name to delete"
              return 1
            fi
          fi
        }

        if [[ $MULTI_ACCOUNT_SCAN = 'true' ]]; then
          echo "Getting list of accounts to scan"
          if [[ $MULTI_ACCOUNT_LIST_OVERRIDE != '' ]]; then
            echo "Using account overrides."
            account_list=${MULTI_ACCOUNT_LIST_OVERRIDE//,/ }
          else
            echo "Using accounts from AWS Organizations."
            if ! account_list=$(aws organizations list-accounts --query 'Accounts[?Status==`ACTIVE`].Id' --output text); then
              echo "ERROR: Failed to retrieve account list from AWS Organizations"
              exit 1
            fi
          fi

          if [[ -z "$account_list" ]]; then
            echo "ERROR: No accounts found to scan"
            exit 1
          fi

          echo "Will scan accounts: $account_list"

          # Persist failures and execution ARNs across CodeBuild phases.
          ASSESSMENT_FAILURES_FILE=/tmp/assessment_failures.tsv
          EXPECTED_ACCOUNTS_FILE=/tmp/expected_accounts.txt
          : > "$ASSESSMENT_FAILURES_FILE"
          : > "$EXPECTED_ACCOUNTS_FILE"
          for accountId in $account_list; do
            printf '%s\n' "$accountId" >> "$EXPECTED_ACCOUNTS_FILE"
          done
          # The management account is always assessed separately below, even
          # when an override lists only member accounts. Persist it here so
          # post-build must poll, validate, and collect that execution too.
          if ! grep -Fxq "$AWS_ACCOUNT_ID" "$EXPECTED_ACCOUNTS_FILE"; then
            printf '%s\n' "$AWS_ACCOUNT_ID" >> "$EXPECTED_ACCOUNTS_FILE"
          fi
          mkdir -p /tmp/execution_arns

          record_failure() {
            local account_id="$1"
            local stage="$2"
            local reason="$3"
            printf '%s\t%s\t%s\n' "$account_id" "$stage" "$reason" >> "$ASSESSMENT_FAILURES_FILE"
          }

          for accountId in $account_list; do
            echo "Processing account $accountId"
            if [[ $accountId == $AWS_ACCOUNT_ID ]]; then
              echo "Skipping management account $accountId - will be handled separately"
              continue
            fi

            # Assume role with error handling - keep credentials in memory only
            CREDS=$(aws sts assume-role --role-arn arn:$AWS_PARTITION:iam::$accountId:role/service-role/$MEMBER_ROLE --role-session-name AIMLSecurityAssessment --output json 2>/dev/null)
            if [[ $? -ne 0 || -z "$CREDS" ]]; then
              echo "ERROR: Failed to assume role in account $accountId"
              record_failure "$accountId" "deployment" "Failed to assume the member role"
              continue
            fi

            # Extract credentials from memory
            export AWS_ACCESS_KEY_ID=$(echo "$CREDS" | jq -r '.Credentials.AccessKeyId')
            export AWS_SECRET_ACCESS_KEY=$(echo "$CREDS" | jq -r '.Credentials.SecretAccessKey')
            export AWS_SESSION_TOKEN=$(echo "$CREDS" | jq -r '.Credentials.SessionToken')
            unset CREDS

            if [[ -z "$AWS_ACCESS_KEY_ID" || "$AWS_ACCESS_KEY_ID" == "null" ]]; then
              echo "ERROR: Failed to extract credentials for account $accountId"
              record_failure "$accountId" "deployment" "AssumeRole returned incomplete credentials"
              unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
              continue
            fi

            echo "Deploying to account $accountId"
            if ! recover_failed_stack aws-sam-cli-managed-default || ! recover_failed_stack "aiml-security-$accountId"; then
              echo "ERROR: Failed to recover deployment stack for account $accountId"
              record_failure "$accountId" "deployment" "Failed to recover a failed CloudFormation stack"
              unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
              continue
            elif ! sam deploy --template-file .aws-sam/build/template.yaml --stack-name "aiml-security-$accountId" --capabilities CAPABILITY_IAM --no-confirm-changeset --no-fail-on-empty-changeset --resolve-s3 --region "$AWS_DEFAULT_REGION" --parameter-overrides "$SAM_TARGET_REGIONS_PARAMETER" "$SAM_ENABLE_BEDROCK_PARAMETER" "$SAM_ENABLE_SAGEMAKER_PARAMETER" "$SAM_ENABLE_AGENTCORE_PARAMETER" "$SAM_ENABLE_AGENT_REGISTRY_PARAMETER" "$SAM_REQUIRE_BEDROCK_ZERO_DATA_RETENTION_PARAMETER" "$SAM_REQUIRE_MARKETPLACE_ENDPOINT_CMK_PARAMETER" "$SAM_REQUIRE_AGENTCORE_ONLINE_EVALUATION_PARAMETER" "$SAM_REQUIRE_AGENT_REGISTRY_CMK_PARAMETER" "$SAM_AGENTCORE_TOKEN_VAULT_ID_PARAMETER" "$SAM_ENABLE_AGENTCORE_ARTIFACT_OBJECT_READS_PARAMETER" "$SAM_APPROVED_EXTERNAL_ACCOUNT_IDS_PARAMETER" "$SAM_APPROVED_ORGANIZATION_IDS_PARAMETER" "$SAM_ENABLE_SAGEMAKER_ARTIFACT_OBJECT_READS_PARAMETER"; then
              echo "ERROR: Deploy failed for account $accountId"
              record_failure "$accountId" "deployment" "SAM deployment failed"
              unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
              continue
            fi

            STATE_MACHINE_ARN=$(aws cloudformation describe-stacks --stack-name aiml-security-$accountId --query 'Stacks[0].Outputs[?OutputKey==`AIMLAssessmentStateMachineArn`].OutputValue' --output text 2>/dev/null)
            if [[ $STATE_MACHINE_ARN != "" && $STATE_MACHINE_ARN != "None" ]]; then
              if EXECUTION_ARN=$(aws stepfunctions start-execution --state-machine-arn $STATE_MACHINE_ARN --input "{\"accountId\":\"$accountId\",\"enableResponsibleAIGRC\":\"$ENABLE_RESPONSIBLE_AI_GRC\",\"enableOWASP\":\"$ENABLE_OWASP\"}" --query 'executionArn' --output text); then
                if [[ -n "$EXECUTION_ARN" && "$EXECUTION_ARN" != "None" ]]; then
                  echo "Started execution: $EXECUTION_ARN"
                  echo "$EXECUTION_ARN" > /tmp/execution_arns/$accountId.txt
                else
                  echo "ERROR: Step Functions returned an empty execution ARN for account $accountId"
                  record_failure "$accountId" "execution-start" "Step Functions returned an empty execution ARN"
                fi
              else
                echo "ERROR: Failed to start Step Function execution for account $accountId"
                record_failure "$accountId" "execution-start" "Failed to start the Step Functions execution"
              fi
            else
              echo "ERROR: No State Machine ARN found for account $accountId"
              record_failure "$accountId" "execution-start" "The deployed stack did not return a state-machine ARN"
            fi

            unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
          done

          echo "Deploying to management account $AWS_ACCOUNT_ID"
          if ! recover_failed_stack aws-sam-cli-managed-default || ! recover_failed_stack aiml-security-mgmt; then
            echo "ERROR: Failed to recover deployment stack for management account"
            record_failure "$AWS_ACCOUNT_ID" "deployment" "Failed to recover a failed CloudFormation stack"
          elif ! sam deploy --template-file .aws-sam/build/template.yaml --stack-name aiml-security-mgmt --capabilities CAPABILITY_IAM --no-confirm-changeset --no-fail-on-empty-changeset --resolve-s3 --region "$AWS_DEFAULT_REGION" --parameter-overrides "$SAM_TARGET_REGIONS_PARAMETER" "$SAM_ENABLE_BEDROCK_PARAMETER" "$SAM_ENABLE_SAGEMAKER_PARAMETER" "$SAM_ENABLE_AGENTCORE_PARAMETER" "$SAM_ENABLE_AGENT_REGISTRY_PARAMETER" "$SAM_REQUIRE_BEDROCK_ZERO_DATA_RETENTION_PARAMETER" "$SAM_REQUIRE_MARKETPLACE_ENDPOINT_CMK_PARAMETER" "$SAM_REQUIRE_AGENTCORE_ONLINE_EVALUATION_PARAMETER" "$SAM_REQUIRE_AGENT_REGISTRY_CMK_PARAMETER" "$SAM_AGENTCORE_TOKEN_VAULT_ID_PARAMETER" "$SAM_ENABLE_AGENTCORE_ARTIFACT_OBJECT_READS_PARAMETER" "$SAM_APPROVED_EXTERNAL_ACCOUNT_IDS_PARAMETER" "$SAM_APPROVED_ORGANIZATION_IDS_PARAMETER" "$SAM_ENABLE_SAGEMAKER_ARTIFACT_OBJECT_READS_PARAMETER"; then
            echo "ERROR: Deploy failed for management account"
            record_failure "$AWS_ACCOUNT_ID" "deployment" "SAM deployment failed"
          else
            STATE_MACHINE_ARN=$(aws cloudformation describe-stacks --stack-name aiml-security-mgmt --query 'Stacks[0].Outputs[?OutputKey==`AIMLAssessmentStateMachineArn`].OutputValue' --output text 2>/dev/null)
            if [[ $STATE_MACHINE_ARN != "" && $STATE_MACHINE_ARN != "None" ]]; then
              if MGMT_EXECUTION_ARN=$(aws stepfunctions start-execution --state-machine-arn $STATE_MACHINE_ARN --input "{\"accountId\":\"$AWS_ACCOUNT_ID\",\"enableResponsibleAIGRC\":\"$ENABLE_RESPONSIBLE_AI_GRC\",\"enableOWASP\":\"$ENABLE_OWASP\"}" --query 'executionArn' --output text); then
                if [[ -n "$MGMT_EXECUTION_ARN" && "$MGMT_EXECUTION_ARN" != "None" ]]; then
                  echo "Started management account execution: $MGMT_EXECUTION_ARN"
                  echo "$MGMT_EXECUTION_ARN" > /tmp/execution_arns/$AWS_ACCOUNT_ID.txt
                else
                  echo "ERROR: Step Functions returned an empty execution ARN for management account"
                  record_failure "$AWS_ACCOUNT_ID" "execution-start" "Step Functions returned an empty execution ARN"
                fi
              else
                echo "ERROR: Failed to start Step Function execution for management account"
                record_failure "$AWS_ACCOUNT_ID" "execution-start" "Failed to start the Step Functions execution"
              fi
            else
              echo "ERROR: No State Machine ARN found for management account"
              record_failure "$AWS_ACCOUNT_ID" "execution-start" "The deployed stack did not return a state-machine ARN"
            fi
          fi

          if [[ -s "$ASSESSMENT_FAILURES_FILE" ]]; then
            echo "One or more accounts could not start an assessment; post-build will process the remaining accounts and fail after reporting all gaps."
          fi

        else
          echo "Single account deployment"
          STACK_NAME="aiml-sec-${AWS_ACCOUNT_ID}"
          if ! recover_failed_stack aws-sam-cli-managed-default || ! recover_failed_stack "$STACK_NAME"; then
            echo "ERROR: Failed to recover deployment stack for single account"
            exit 1
          elif ! sam deploy --template-file .aws-sam/build/template.yaml --stack-name "$STACK_NAME" --capabilities CAPABILITY_IAM --no-confirm-changeset --no-fail-on-empty-changeset --resolve-s3 --region "$AWS_DEFAULT_REGION" --parameter-overrides "$SAM_TARGET_REGIONS_PARAMETER" "$SAM_ENABLE_BEDROCK_PARAMETER" "$SAM_ENABLE_SAGEMAKER_PARAMETER" "$SAM_ENABLE_AGENTCORE_PARAMETER" "$SAM_ENABLE_AGENT_REGISTRY_PARAMETER" "$SAM_REQUIRE_BEDROCK_ZERO_DATA_RETENTION_PARAMETER" "$SAM_REQUIRE_MARKETPLACE_ENDPOINT_CMK_PARAMETER" "$SAM_REQUIRE_AGENTCORE_ONLINE_EVALUATION_PARAMETER" "$SAM_REQUIRE_AGENT_REGISTRY_CMK_PARAMETER" "$SAM_AGENTCORE_TOKEN_VAULT_ID_PARAMETER" "$SAM_ENABLE_AGENTCORE_ARTIFACT_OBJECT_READS_PARAMETER" "$SAM_APPROVED_EXTERNAL_ACCOUNT_IDS_PARAMETER" "$SAM_APPROVED_ORGANIZATION_IDS_PARAMETER" "$SAM_ENABLE_SAGEMAKER_ARTIFACT_OBJECT_READS_PARAMETER"; then
            echo "ERROR: SAM deploy failed for single account"
            exit 1
          fi
          STATE_MACHINE_ARN=$(aws cloudformation describe-stacks --stack-name $STACK_NAME --query 'Stacks[0].Outputs[?OutputKey==`AIMLAssessmentStateMachineArn`].OutputValue' --output text)
          if [[ -z "$STATE_MACHINE_ARN" || "$STATE_MACHINE_ARN" == "None" ]]; then
            echo "ERROR: Failed to retrieve State Machine ARN"
            exit 1
          fi
          # Capture execution ARN for post-build phase to wait on
          EXECUTION_ARN=$(aws stepfunctions start-execution --state-machine-arn $STATE_MACHINE_ARN --input "{\"accountId\":\"$AWS_ACCOUNT_ID\",\"enableResponsibleAIGRC\":\"$ENABLE_RESPONSIBLE_AI_GRC\",\"enableOWASP\":\"$ENABLE_OWASP\"}" --query 'executionArn' --output text)
          if [[ -z "$EXECUTION_ARN" || "$EXECUTION_ARN" == "None" ]]; then
            echo "ERROR: Failed to start Step Function execution"
            exit 1
          fi
          echo "Started Step Function execution: $EXECUTION_ARN"
          # Save execution ARN to file for post-build phase
          echo "$EXECUTION_ARN" > /tmp/single_account_execution_arn.txt
        fi
  post_build:
    commands:
      - |
        # Older top-level stacks do not define these project variables.
        # Reapply defaults in this phase so the artifact gate stays complete.
        export ENABLE_BEDROCK="${ENABLE_BEDROCK:-true}"
        export ENABLE_SAGEMAKER="${ENABLE_SAGEMAKER:-true}"
        export ENABLE_AGENTCORE="${ENABLE_AGENTCORE:-true}"
        export ENABLE_AGENT_REGISTRY="${ENABLE_AGENT_REGISTRY:-true}"
      - echo "CODEBUILD_BUILD_SUCCEEDING status is $CODEBUILD_BUILD_SUCCEEDING"
      - |
        if [ $CODEBUILD_BUILD_SUCCEEDING -eq 0 ]; then
          echo "ERROR: The previous build stage failed. Exiting post_build without processing results."
          exit 1
        fi

        echo "Assessment completed. Results in S3:$BUCKET_REPORT"

        # When Responsible AI GRC runs only as an OWASP dependency
        # (EnableResponsibleAIGRCAssessment=false but EnableOWASPAssessment=true),
        # its CSV must never reach the customer-facing report bucket. The report
        # layer already strips FS-* rows from the HTML, but the raw
        # responsible_ai_grc_security_report_*.csv would otherwise be copied
        # verbatim. Build an exclude flag (empty when Responsible AI GRC is
        # enabled) and apply it to every copy/sync into the report bucket.
        # ENABLE_RESPONSIBLE_AI_GRC is already normalised to lowercase in the
        # build phase, so a direct "true" comparison is sufficient here.
        fs_exclude_args=()
        if [[ "${ENABLE_RESPONSIBLE_AI_GRC:-false}" != "true" ]]; then
          fs_exclude_args=(--exclude "responsible_ai_grc_security_report_*.csv")
        fi

        # "Changes since last assessment" report (assessment_history package,
        # EnableAssessmentHistory setting). Called once in each branch below,
        # after the existing reports. It can never fail the build: problems are
        # printed as WARNING lines, the tool handles each account separately and
        # exits 0, and the step has a time limit. It only reads the failure
        # ledger and never writes to it.
        run_changes_report() {
          local enabled
          enabled="$(echo "${ENABLE_ASSESSMENT_HISTORY:-true}" | tr '[:upper:]' '[:lower:]')"
          if [[ "$enabled" == "false" ]]; then
            echo "Changes report disabled (EnableAssessmentHistory=false)"
            return 0
          fi
          if [[ "$enabled" != "true" ]]; then
            echo "WARNING: ENABLE_ASSESSMENT_HISTORY is '$ENABLE_ASSESSMENT_HISTORY', not true or false; using true"
          fi
          # Time budget: one minute per account, at least five minutes, so
          # large organizations aren't cut short (review item F4). The tool
          # stops starting accounts a minute before the budget ends; `timeout`
          # stops it 30 seconds after (review item L2).
          local arg previous_arg="" accounts=""
          for arg in "$@"; do
            if [[ "$previous_arg" == "--accounts" ]]; then
              accounts="$arg"
            fi
            previous_arg="$arg"
          done
          local account_count
          account_count=$(echo "$accounts" | tr ',' ' ' | wc -w)
          local budget=$(( account_count * 60 ))
          if (( budget < 300 )); then
            budget=300
          fi
          # Skip when little build time is left, and never budget more than is
          # left, so this step can't push the build past its CodeBuild timeout
          # (CODEBUILD_START_TIME is in ms).
          if [[ -n "${CODEBUILD_START_TIME:-}" ]]; then
            local seconds_left=$(( ${BUILD_TIMEOUT_MINUTES:-60} * 60 - ( $(date +%s) - CODEBUILD_START_TIME / 1000 ) ))
            if (( seconds_left < 360 )); then
              echo "WARNING: Changes report skipped: only ${seconds_left}s of build time left"
              return 0
            fi
            if (( budget > seconds_left - 60 )); then
              budget=$(( seconds_left - 60 ))
            fi
          fi
          echo "Creating changes-since-last-assessment reports..."
          if ! (cd .. && ENABLE_ASSESSMENT_HISTORY=true timeout $(( budget + 30 )) python3 -m assessment_history compare --bucket "$BUCKET_REPORT" --time-budget "$budget" "$@"); then
            echo "WARNING: The changes report step did not complete; assessment results are not affected"
          fi
          return 0
        }

        # Single-account mode: record whether this run's Step Functions
        # execution succeeded, next to its CSVs in the central bucket, so the
        # changes report never uses a failed run as its previous run. An
        # execution only succeeds when the main report's completeness check
        # passes. The record also lists the core services the run selected,
        # so a later run can tell a deselected service from a missing CSV.
        # Written only while the setting is on; it can never fail the build.
        record_assessment_run() {
          local execution_id="${1##*:}" status="${2:-UNKNOWN}" succeeded=true
          if [[ "$3" == true ]]; then
            succeeded=false
          fi
          if [[ "$(echo "${ENABLE_ASSESSMENT_HISTORY:-true}" | tr '[:upper:]' '[:lower:]')" == "false" ]]; then
            return 0
          fi
          if [[ -z "$execution_id" || "$execution_id" == "None" ]]; then
            echo "WARNING: No execution ID was saved, so no assessment run record was written"
            return 0
          fi
          local selection="" pair flag value
          for pair in bedrock:ENABLE_BEDROCK sagemaker:ENABLE_SAGEMAKER agentcore:ENABLE_AGENTCORE agent-registry:ENABLE_AGENT_REGISTRY; do
            flag="${pair#*:}"
            value="${!flag:-true}"
            if [[ "$value" != true && "$value" != false ]]; then
              echo "WARNING: $flag is '$value', not true or false, so the run record does not list the selected services"
              selection=""
              break
            fi
            selection+="${selection:+, }\"${pair%%:*}\": $value"
          done
          if [[ -n "$selection" ]]; then
            selection=", \"selected_services\": {$selection}"
          fi
          local key="$AWS_ACCOUNT_ID/assessment_history_run_${execution_id}.json"
          if ! printf '{"execution_id": "%s", "succeeded": %s, "status": "%s", "written_at": "%s"%s}\n' \
              "$execution_id" "$succeeded" "$status" "$(date -u +%Y-%m-%dT%H:%M:%SZ)" "$selection" \
              | aws s3 cp - "s3://$BUCKET_REPORT/$key" --content-type application/json --only-show-errors; then
            echo "WARNING: Could not write the assessment run record s3://$BUCKET_REPORT/$key"
          fi
          return 0
        }

        # Derive how long to wait for Step Function completion from the overall
        # CodeBuild timeout, reserving a buffer for the build/deploy work that runs
        # before this loop and the report processing that runs after it. An explicit
        # SF_POLL_TIMEOUT override always wins; a floor guards against tiny/negative
        # values when BUILD_TIMEOUT_MINUTES is small.
        SF_POLL_BUFFER_SECONDS=${SF_POLL_BUFFER_SECONDS:-600}
        SF_POLL_FLOOR_SECONDS=${SF_POLL_FLOOR_SECONDS:-300}
        if [[ -n "$SF_POLL_TIMEOUT" ]]; then
          sf_poll_timeout=$SF_POLL_TIMEOUT
        else
          sf_poll_timeout=$(( ${BUILD_TIMEOUT_MINUTES:-60} * 60 - SF_POLL_BUFFER_SECONDS ))
          if [[ $sf_poll_timeout -lt $SF_POLL_FLOOR_SECONDS ]]; then
            sf_poll_timeout=$SF_POLL_FLOOR_SECONDS
          fi
        fi
        echo "Step Function poll timeout set to ${sf_poll_timeout}s (CodeBuild timeout: ${BUILD_TIMEOUT_MINUTES:-60}m, buffer: ${SF_POLL_BUFFER_SECONDS}s, floor: ${SF_POLL_FLOOR_SECONDS}s)"

        if [[ $MULTI_ACCOUNT_SCAN = 'true' ]]; then
          EXPECTED_ACCOUNTS_FILE=/tmp/expected_accounts.txt
          if [[ ! -s "$EXPECTED_ACCOUNTS_FILE" ]]; then
            echo "ERROR: Expected account list was not persisted by the build phase"
            exit 1
          fi
          account_list=$(tr '\n' ' ' < "$EXPECTED_ACCOUNTS_FILE")
          echo "Account list for post-build processing: $account_list"

          # Continue the failure ledger created during the build phase.
          ASSESSMENT_FAILURES_FILE=/tmp/assessment_failures.tsv
          touch "$ASSESSMENT_FAILURES_FILE"

          record_failure() {
            local account_id="$1"
            local stage="$2"
            local reason="$3"
            printf '%s\t%s\t%s\n' "$account_id" "$stage" "$reason" >> "$ASSESSMENT_FAILURES_FILE"
          }

          # Clean up any existing account-files directory to ensure fresh start
          echo "Cleaning up existing account-files directory"
          rm -rf /tmp/account-files
          mkdir -p /tmp/account-files

          echo "Copying files from respective account S3 buckets to CodeBuild environment"
          for accountId in $account_list; do
            echo "Processing account $accountId"
            if [[ $accountId == $AWS_ACCOUNT_ID ]]; then
              # Management account - use mgmt stack name
              STACK_NAME="aiml-security-mgmt"
            else
              # Member account - assume role and use account-specific stack name
              CREDS=$(aws sts assume-role --role-arn arn:$AWS_PARTITION:iam::$accountId:role/service-role/$MEMBER_ROLE --role-session-name AIMLSecurityAssessment --output json 2>/dev/null)
              if [[ $? -ne 0 || -z "$CREDS" ]]; then
                echo "ERROR: Failed to assume role for account $accountId in post-build"
                record_failure "$accountId" "result-collection" "Failed to assume the member role"
                continue
              fi

              export AWS_ACCESS_KEY_ID=$(echo "$CREDS" | jq -r '.Credentials.AccessKeyId')
              export AWS_SECRET_ACCESS_KEY=$(echo "$CREDS" | jq -r '.Credentials.SecretAccessKey')
              export AWS_SESSION_TOKEN=$(echo "$CREDS" | jq -r '.Credentials.SessionToken')
              unset CREDS

              if [[ -z "$AWS_ACCESS_KEY_ID" || "$AWS_ACCESS_KEY_ID" == "null" ]]; then
                echo "ERROR: Failed to extract credentials for account $accountId in post-build"
                record_failure "$accountId" "result-collection" "AssumeRole returned incomplete credentials"
                unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
                continue
              fi

              STACK_NAME="aiml-security-$accountId"
            fi

            execution_succeeded=false

            # Wait for Step Functions to complete before copying files.
            if [[ -f /tmp/execution_arns/$accountId.txt ]]; then
              EXECUTION_ARN=$(cat /tmp/execution_arns/$accountId.txt)
              echo "Waiting for Step Function execution in account $accountId: $EXECUTION_ARN"
              if [[ $EXECUTION_ARN != "" && $EXECUTION_ARN != "None" ]]; then
                timeout=$sf_poll_timeout
                elapsed=0
                while [[ $elapsed -lt $timeout ]]; do
                  if ! STATUS=$(aws stepfunctions describe-execution --execution-arn "$EXECUTION_ARN" --query 'status' --output text 2>/dev/null); then
                    echo "ERROR: Failed to poll Step Functions execution for account $accountId"
                    record_failure "$accountId" "execution" "Failed to query the Step Functions execution"
                    break
                  fi
                  case "$STATUS" in
                    SUCCEEDED)
                      echo "Step Function for account $accountId completed successfully"
                      execution_succeeded=true
                      break
                      ;;
                    FAILED|TIMED_OUT|ABORTED|PENDING_REDRIVE)
                      echo "ERROR: Step Function did not succeed for account $accountId (status: $STATUS)"
                      record_failure "$accountId" "execution" "Step Functions completed with status $STATUS"
                      break
                      ;;
                    RUNNING)
                      ;;
                    *)
                      echo "ERROR: Step Functions returned unexpected status '$STATUS' for account $accountId"
                      record_failure "$accountId" "execution" "Step Functions returned unexpected status $STATUS"
                      break
                      ;;
                  esac
                  echo "Step Function for account $accountId still running... waiting 30 seconds"
                  sleep 30
                  elapsed=$((elapsed + 30))
                done
                if [[ $execution_succeeded != true && $elapsed -ge $timeout ]]; then
                  echo "ERROR: Timeout waiting for Step Function completion in account $accountId"
                  record_failure "$accountId" "execution" "Timed out waiting for Step Functions completion"
                fi
              else
                echo "ERROR: Invalid execution ARN for account $accountId"
                record_failure "$accountId" "execution" "The saved execution ARN was empty"
              fi
            else
              echo "ERROR: No execution ARN file found for account $accountId"
              record_failure "$accountId" "execution" "No execution ARN was saved during deployment"
            fi

            if [[ $execution_succeeded != true ]]; then
              if [[ $accountId != $AWS_ACCOUNT_ID ]]; then
                unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
              fi
              continue
            fi

            if ! ACCOUNT_BUCKET=$(aws cloudformation describe-stacks --stack-name "$STACK_NAME" --query 'Stacks[0].Outputs[?OutputKey==`AssessmentBucketName`].OutputValue' --output text 2>/dev/null); then
              ACCOUNT_BUCKET=""
            fi
            if [[ $ACCOUNT_BUCKET != "" && $ACCOUNT_BUCKET != "None" ]]; then
              echo "Copying files from $ACCOUNT_BUCKET to local storage"
              mkdir -p /tmp/account-files/$accountId
              execution_id="${EXECUTION_ARN##*:}"

              if ! aws s3 cp "s3://$ACCOUNT_BUCKET/" "/tmp/account-files/$accountId/" --recursive --exclude "*" --include "*_security_report_${execution_id}*.csv" --include "security_assessment_single_account_*.html" --exclude "*/" "${fs_exclude_args[@]}"; then
                echo "ERROR: Failed to copy assessment artifacts for account $accountId"
                record_failure "$accountId" "artifact-copy" "Failed to copy assessment artifacts from the account bucket"
                rm -rf "/tmp/account-files/$accountId"
              else
                required_artifact_prefixes=()
                if [[ $ENABLE_BEDROCK == "true" ]]; then
                  required_artifact_prefixes+=(bedrock)
                fi
                if [[ $ENABLE_SAGEMAKER == "true" ]]; then
                  required_artifact_prefixes+=(sagemaker)
                fi
                if [[ $ENABLE_AGENTCORE == "true" ]]; then
                  required_artifact_prefixes+=(agentcore)
                fi
                if [[ $ENABLE_AGENT_REGISTRY == "true" ]]; then
                  required_artifact_prefixes+=(agent_registry)
                fi
                if [[ $ENABLE_RESPONSIBLE_AI_GRC == "true" ]]; then
                  required_artifact_prefixes+=(responsible_ai_grc)
                fi
                if [[ $ENABLE_OWASP == "true" ]]; then
                  required_artifact_prefixes+=(owasp)
                fi

                artifacts_complete=true
                for artifact_prefix in "${required_artifact_prefixes[@]}"; do
                  artifact_match=$(find "/tmp/account-files/$accountId" -maxdepth 1 -type f -name "${artifact_prefix}_security_report_${execution_id}*.csv" -print -quit)
                  if [[ -z "$artifact_match" ]]; then
                    echo "ERROR: Missing $artifact_prefix CSV for account $accountId and execution $execution_id"
                    record_failure "$accountId" "artifact-validation" "Missing $artifact_prefix CSV for execution $execution_id"
                    artifacts_complete=false
                  fi
                done

                html_match=$(find "/tmp/account-files/$accountId" -maxdepth 1 -type f -name "security_assessment_single_account_*.html" -print -quit)
                if [[ -z "$html_match" ]]; then
                  echo "ERROR: Missing HTML report for account $accountId"
                  record_failure "$accountId" "artifact-validation" "Missing the per-account HTML report"
                  artifacts_complete=false
                fi

                if [[ $artifacts_complete != true ]]; then
                  rm -rf "/tmp/account-files/$accountId"
                fi
              fi
            else
              echo "ERROR: No assessment bucket found for $accountId (STACK_NAME: $STACK_NAME)"
              record_failure "$accountId" "result-collection" "The deployed stack did not return an assessment bucket"
            fi

            if [[ $accountId != $AWS_ACCOUNT_ID ]]; then
              unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
            fi
          done

          echo "Uploading files to management account S3 bucket and creating consolidated report"

          # Debug: Show which directories exist
          echo "Available account directories:"
          ls -la /tmp/account-files/ || echo "No account-files directory"

          # Upload files from local storage to management account bucket organized by account (overwrite)
          echo "Uploading account reports to S3..."
          for accountId in $account_list; do
            if [[ -d /tmp/account-files/$accountId ]]; then
              echo "Uploading files for account $accountId (overwriting existing)"
              if ! aws s3 sync /tmp/account-files/$accountId/ s3://$BUCKET_REPORT/$accountId/ "${fs_exclude_args[@]}"; then
                echo "ERROR: Failed to upload files for $accountId"
                record_failure "$accountId" "artifact-upload" "Failed to upload account artifacts to the central report bucket"
              fi
            else
              echo "ERROR: No validated artifact directory found for account $accountId"
              record_failure "$accountId" "artifact-validation" "No complete artifact set was available for consolidation"
            fi
          done

          if [[ -s "$ASSESSMENT_FAILURES_FILE" ]]; then
            echo "Skipping the consolidated report because one or more accounts have incomplete coverage."
          else
            echo "Creating consolidated HTML report..."
            if ! python3 ../consolidate_html_reports.py; then
              echo "ERROR: Failed to create consolidated HTML report"
              record_failure "$AWS_ACCOUNT_ID" "consolidation" "Failed to create or upload the consolidated HTML report"
            fi
          fi

          run_changes_report --accounts "$account_list" --execution-arn-dir /tmp/execution_arns --failures-file "$ASSESSMENT_FAILURES_FILE"

          if [[ -s "$ASSESSMENT_FAILURES_FILE" ]]; then
            echo "==========================================="
            echo "Multi-account assessment incomplete"
            echo "Account | Stage | Reason"
            while IFS=$'\t' read -r failed_account failed_stage failed_reason; do
              printf '%s | %s | %s\n' "$failed_account" "$failed_stage" "$failed_reason"
            done < "$ASSESSMENT_FAILURES_FILE"
            echo "==========================================="
            exit 1
          fi

          echo "==========================================="
          echo "Multi-account assessment completed successfully!"
          echo "==========================================="
          echo "ACCESS YOUR RESULTS HERE: s3://$BUCKET_REPORT/"
          echo ""
          echo "Note: Other S3 buckets created during deployment (e.g., aiml-sec-*-aimlassessmentbucket-*)"
          echo "are for internal use. Use the bucket above from the AssessmentBucket stack output."
          echo "==========================================="
        else
          # Single account post-build: wait for Step Function and report status
          echo "Single account post-build processing"
          STACK_NAME="aiml-sec-${AWS_ACCOUNT_ID}"
          sf_failed=false

          # Read the execution ARN saved from build phase
          if [[ -f /tmp/single_account_execution_arn.txt ]]; then
            EXECUTION_ARN=$(cat /tmp/single_account_execution_arn.txt)
            echo "Waiting for Step Function execution: $EXECUTION_ARN"
            if [[ $EXECUTION_ARN != "" && $EXECUTION_ARN != "None" ]]; then
              timeout=$sf_poll_timeout
              elapsed=0
              while [[ $elapsed -lt $timeout ]]; do
                STATUS=$(aws stepfunctions describe-execution --execution-arn $EXECUTION_ARN --query 'status' --output text 2>/dev/null)
                if [[ $STATUS == "SUCCEEDED" || $STATUS == "FAILED" || $STATUS == "TIMED_OUT" || $STATUS == "ABORTED" ]]; then
                  echo "Step Function completed with status: $STATUS"
                  if [[ $STATUS != "SUCCEEDED" ]]; then
                    echo "ERROR: Step Function execution failed with status: $STATUS"
                    sf_failed=true
                  fi
                  break
                fi
                echo "Step Function still running (status: $STATUS)... waiting 30 seconds"
                sleep 30
                elapsed=$((elapsed + 30))
              done
              if [[ $elapsed -ge $timeout ]]; then
                echo "ERROR: Timeout waiting for Step Function completion"
                sf_failed=true
              fi
            fi
          else
            echo "ERROR: No execution ARN file found from build phase"
            sf_failed=true
          fi

          # Get assessment bucket and sync results to BUCKET_REPORT for consistency
          ACCOUNT_BUCKET=$(aws cloudformation describe-stacks --stack-name $STACK_NAME --query 'Stacks[0].Outputs[?OutputKey==`AssessmentBucketName`].OutputValue' --output text 2>/dev/null)
          if [[ $ACCOUNT_BUCKET != "" && $ACCOUNT_BUCKET != "None" ]]; then
            echo "Assessment results generated in S3 bucket: $ACCOUNT_BUCKET"
            echo "Bucket contents:"
            aws s3 ls s3://$ACCOUNT_BUCKET/ --recursive || echo "Failed to list bucket contents"

            # Sync results to BUCKET_REPORT for consistent access
            if [[ $ACCOUNT_BUCKET != $BUCKET_REPORT ]]; then
              echo "Syncing results to consolidated bucket: $BUCKET_REPORT"
              if ! aws s3 sync s3://$ACCOUNT_BUCKET/ s3://$BUCKET_REPORT/$AWS_ACCOUNT_ID/ --exclude "*" --include "*.csv" --include "*.html" "${fs_exclude_args[@]}"; then
                echo "WARNING: Failed to sync results to $BUCKET_REPORT"
              else
                echo "Results synced successfully to s3://$BUCKET_REPORT/$AWS_ACCOUNT_ID/"
              fi
            else
              echo "Assessment bucket and report bucket are the same, no sync needed"
            fi
          else
            echo "WARNING: Could not determine assessment bucket from stack outputs"
          fi

          record_assessment_run "${EXECUTION_ARN:-}" "${STATUS:-}" "$sf_failed"

          if [[ $sf_failed == true ]]; then
            echo "Changes report skipped: the assessment run did not succeed"
          else
            run_changes_report --accounts "$AWS_ACCOUNT_ID" --execution-id "${EXECUTION_ARN##*:}"
          fi

          if [[ $sf_failed == true ]]; then
            echo "ERROR: Assessment failed. Check Step Functions execution logs for details."
            exit 1
          fi
        fi
      - |
        echo "==========================================="
        echo "Assessment completed successfully!"
        echo "==========================================="
        echo "ACCESS YOUR RESULTS HERE: s3://$BUCKET_REPORT/"
        echo ""
        echo "Note: Other S3 buckets created during deployment (e.g., aiml-sec-*-aimlassessmentbucket-*)"
        echo "are for internal use. Use the bucket above from the AssessmentBucket stack output."
        echo "==========================================="