buildspec.yml

version: 0.2
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"
      - |
        # 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
      - 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/
      - |
        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"

          # Track failed accounts
          failed_accounts=()

          # Create directory to store execution ARNs for post-build phase
          mkdir -p /tmp/execution_arns

          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"
              failed_accounts+=($accountId)
              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"
              failed_accounts+=($accountId)
              unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
              continue
            fi

            echo "Deploying to account $accountId"
            if ! 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; then
              echo "ERROR: Deploy failed for account $accountId"
              failed_accounts+=($accountId)
              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\"}" --query 'executionArn' --output text); then
                echo "Started execution: $EXECUTION_ARN"
                # Save execution ARN for post-build phase
                echo "$EXECUTION_ARN" > /tmp/execution_arns/$accountId.txt
              else
                echo "WARNING: Failed to start Step Function execution for account $accountId"
              fi
            else
              echo "WARNING: No State Machine ARN found for account $accountId"
            fi

            unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
          done

          echo "Deploying to management account $AWS_ACCOUNT_ID"
          # Clean up any leftover SAM managed stack from previous failed runs
          SAM_STATUS=$(aws cloudformation describe-stacks --stack-name aws-sam-cli-managed-default --query 'Stacks[0].StackStatus' --output text 2>/dev/null)
          if [[ "$SAM_STATUS" == "ROLLBACK_COMPLETE" || "$SAM_STATUS" == "DELETE_FAILED" ]]; then
            echo "Cleaning up leftover SAM managed stack (status: $SAM_STATUS)"
            aws cloudformation delete-stack --stack-name aws-sam-cli-managed-default
            aws cloudformation wait stack-delete-complete --stack-name aws-sam-cli-managed-default 2>/dev/null || true
          fi
          if ! 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; then
            echo "ERROR: Deploy failed for management account"
            failed_accounts+=($AWS_ACCOUNT_ID)
          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\"}" --query 'executionArn' --output text); then
                echo "Started management account execution: $MGMT_EXECUTION_ARN"
                # Save execution ARN for post-build phase
                echo "$MGMT_EXECUTION_ARN" > /tmp/execution_arns/$AWS_ACCOUNT_ID.txt
              else
                echo "WARNING: Failed to start Step Function execution for management account"
              fi
            fi
          fi

          # Report failed accounts
          if [[ ${#failed_accounts[@]} -gt 0 ]]; then
            echo "WARNING: Deployments failed for the following accounts: ${failed_accounts[*]}"
            echo "Check logs above for specific error details"
          fi

        else
          echo "Single account deployment"
          STACK_NAME="aiml-sec-${AWS_ACCOUNT_ID}"
          if ! 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; 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\"}" --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:
      - 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"

        if [[ $MULTI_ACCOUNT_SCAN = 'true' ]]; then
          echo "Getting list of accounts for post-build processing"
          if [[ $MULTI_ACCOUNT_LIST_OVERRIDE != '' ]]; then
            account_list=$MULTI_ACCOUNT_LIST_OVERRIDE
          else
            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
          echo "Account list for post-build processing: $account_list"

          # Track failures
          upload_failed=false
          failed_accounts=()

          # 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"
                failed_accounts+=($accountId)
                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"
                failed_accounts+=($accountId)
                continue
              fi

              STACK_NAME="aiml-security-$accountId"
            fi

            # Wait for Step Function to complete before copying files
            # Read the execution ARN saved from build phase
            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:-600}
                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 for account $accountId completed with status: $STATUS"
                    if [[ $STATUS == "FAILED" || $STATUS == "TIMED_OUT" || $STATUS == "ABORTED" ]]; then
                      echo "WARNING: Step Function did not succeed for account $accountId"
                    fi
                    break
                  fi
                  echo "Step Function for account $accountId still running (status: $STATUS)... waiting 30 seconds"
                  sleep 30
                  elapsed=$((elapsed + 30))
                done
                if [[ $elapsed -ge $timeout ]]; then
                  echo "WARNING: Timeout waiting for Step Function completion in account $accountId"
                fi
              fi
            else
              echo "WARNING: No execution ARN file found for account $accountId - execution may not have started"
            fi

            # Now copy files from the completed assessment
            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 "Copying files from $ACCOUNT_BUCKET to local storage"
              mkdir -p /tmp/account-files/$accountId

              # List bucket contents first for debugging
              echo "Bucket contents for $accountId:"
              if ! aws s3 ls s3://$ACCOUNT_BUCKET/; then
                echo "WARNING: Failed to list bucket contents for $accountId"
              fi

              if ! aws s3 cp s3://$ACCOUNT_BUCKET/ /tmp/account-files/$accountId/ --recursive --exclude "*" --include "*.csv" --include "*.html" --exclude "*/"; then
                echo "WARNING: No files to copy or copy failed from $accountId"
              fi

              # Show what was actually copied
              echo "Files copied for $accountId:"
              ls -la /tmp/account-files/$accountId/ 2>/dev/null || echo "No files in directory"

              # Flatten any nested directories
              find /tmp/account-files/$accountId -type f \( -name "*.csv" -o -name "*.html" \) -exec mv {} /tmp/account-files/$accountId/ \; 2>/dev/null || true
              find /tmp/account-files/$accountId -type d -empty -delete 2>/dev/null || true
            else
              echo "WARNING: No assessment bucket found for $accountId (STACK_NAME: $STACK_NAME)"
            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/; then
                echo "ERROR: Failed to upload files for $accountId"
                upload_failed=true
              fi
            else
              echo "WARNING: No directory found for account $accountId"
            fi
          done

          # Create consolidated HTML report using separate script
          echo "Creating consolidated HTML report..."
          if ! python3 ../consolidate_html_reports.py; then
            echo "ERROR: Failed to create consolidated HTML report"
            upload_failed=true
          fi

          # Report failures
          if [[ ${#failed_accounts[@]} -gt 0 ]]; then
            echo "WARNING: Post-build processing failed for the following accounts: ${failed_accounts[*]}"
          fi

          if [[ $upload_failed == true ]]; then
            echo "ERROR: One or more uploads or report generation failed. Check logs above for details."
            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}"

          # 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:-600}
              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 "WARNING: Step Function did not succeed"
                  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 "WARNING: Timeout waiting for Step Function completion"
              fi
            fi
          else
            echo "WARNING: No execution ARN file found from build phase"
          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"; 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
        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 "==========================================="