#!/bin/bash

# ==========================================================================
# SCRIPT: ec2_schedule_power_management.sh
# DESCRIPTION: Sets up automated power management for EC2 instances based on tags
#              and a schedule using AWS CLI. This script creates a Lambda function
#              and two EventBridge (CloudWatch Events) rules to stop and start
#              EC2 instances at specified times, helping to optimize costs.
#
# USE CASE SCENARIO:
# A development team uses EC2 instances during business hours but wants to
# automatically stop them overnight and on weekends to save costs. This script
# automates the setup of the necessary AWS resources to achieve this.
#
# PREREQUISITES:
# 1.  **AWS CLI:** The AWS Command Line Interface must be installed and configured
#     with credentials that have the necessary permissions.
# 2.  **IAM Permissions:** The principal executing this script must have permissions to:
#     - `iam:CreateRole`, `iam:AttachRolePolicy`, `iam:PutRolePolicy`, `iam:GetRole`
#     - `lambda:CreateFunction`, `lambda:UpdateFunctionCode`, `lambda:UpdateFunctionConfiguration`, `lambda:GetFunction`, `lambda:AddPermission`
#     - `events:PutRule`, `events:PutTargets`, `events:DescribeRule`
#     - `ec2:StartInstances`, `ec2:StopInstances`, `ec2:DescribeInstances` (for the Lambda role)
# 3.  **Existing Resources:**
#     - EC2 instances tagged with the specified `INSTANCE_TAG_KEY` and `INSTANCE_TAG_VALUE`.
#
# HOW TO USE:
# 1.  **Save the script:** Save this content as `ec2_schedule_power_management.sh`.
# 2.  **Make it executable:** `chmod +x ec2_schedule_power_management.sh`
# 3.  **Configure variables:** Open the script and update the `--- Configuration Variables ---
#     section with your specific environment details.
# 4.  **Run from your terminal:** `./ec2_schedule_power_management.sh`
#
# IMPORTANT CONSIDERATIONS:
# - The Lambda function created will have permissions to start/stop *any* EC2 instance.
#   For production, refine the IAM policy to restrict actions to specific resources (e.g., by tags).
# - Cron expressions are in UTC. Adjust `STOP_CRON_EXPRESSION` and `START_CRON_EXPRESSION`
#   according to your desired local time and timezone conversion.
# - This script creates/updates resources. Running it multiple times will update existing resources.
# ==========================================================================

# --- Configuration Variables (REPLACE with your actual values) ---
INSTANCE_TAG_KEY="AutoStop"             # Tag key to identify target EC2 instances
INSTANCE_TAG_VALUE="true"             # Tag value to identify target EC2 instances
STOP_CRON_EXPRESSION="cron(0 19 ? * MON-FRI *)" # Cron expression for stopping instances (e.g., Mon-Fri at 7 PM UTC)
START_CRON_EXPRESSION="cron(0 8 ? * MON-FRI *)" # Cron expression for starting instances (e.g., Mon-Fri at 8 AM UTC)
AWS_REGION="us-east-1"                    # AWS region
# ----------------------------------------------------------------

# Derived names
LAMBDA_FUNCTION_NAME="ec2-power-manager-${INSTANCE_TAG_KEY}-${INSTANCE_TAG_VALUE}"
LAMBDA_ROLE_NAME="lambda-ec2-power-manager-role-${INSTANCE_TAG_KEY}-${INSTANCE_TAG_VALUE}"
EC2_POLICY_NAME="ec2-power-manager-policy-${INSTANCE_TAG_KEY}-${INSTANCE_TAG_VALUE}"
STOP_RULE_NAME="ec2-stop-rule-${INSTANCE_TAG_KEY}-${INSTANCE_TAG_VALUE}"
START_RULE_NAME="ec2-start-rule-${INSTANCE_TAG_KEY}-${INSTANCE_TAG_VALUE}"

# Lambda function code (Python)
read -r -d '' LAMBDA_CODE << EOM
import boto3
import os

def lambda_handler(event, context):
    ec2 = boto3.client('ec2', region_name=os.environ['AWS_REGION'])
    action = event['action']
    tag_key = os.environ['INSTANCE_TAG_KEY']
    tag_value = os.environ['INSTANCE_TAG_VALUE']

    print(f"Received event to {action} EC2 instances with tag {tag_key}={tag_value}")

    filters = [
        {'Name': f'tag:{tag_key}', 'Values': [tag_value]},
        {'Name': 'instance-state-name', 'Values': ['running', 'stopped']}
    ]

    instances = ec2.describe_instances(Filters=filters)['Reservations']
    instance_ids = []
    for reservation in instances:
        for instance in reservation['Instances']:
            instance_ids.append(instance['InstanceId'])

    if not instance_ids:
        print("No instances found matching the criteria.")
        return

    if action == 'stop':
        print(f"Stopping instances: {instance_ids}")
        ec2.stop_instances(InstanceIds=instance_ids)
    elif action == 'start':
        print(f"Starting instances: {instance_ids}")
        ec2.start_instances(InstanceIds=instance_ids)
    else:
        print(f"Invalid action: {action}")

    return {'statusCode': 200, 'body': f'{action} initiated for {len(instance_ids)} instances.'}
EOM

# Create a temporary zip file for the Lambda function
zip_file="/tmp/${LAMBDA_FUNCTION_NAME}.zip"
echo "${LAMBDA_CODE}" > /tmp/index.py
zip -j "${zip_file}" /tmp/index.py > /dev/null

echo "Starting EC2 power management schedule setup in region ${AWS_REGION}..."

# ==========================================================================
# STEP 1: Create/Update IAM Role for the Lambda function.
# This role grants the Lambda function permissions to interact with EC2 and CloudWatch Logs.
# ==========================================================================
echo ">>> Step 1: Creating/Updating IAM Role for Lambda function..."

# Define the trust policy for the Lambda role, allowing Lambda to assume this role.
LAMBDA_TRUST_POLICY='{"Version": "2012-10-17","Statement": [{"Effect": "Allow","Principal": {"Service": "lambda.amazonaws.com"},"Action": "sts:AssumeRole"}]}'

# Create the IAM role. If it exists, `create-role` will fail, so we check and get its ARN.
LAMBDA_ROLE_ARN=$(aws iam create-role \
    --role-name "${LAMBDA_ROLE_NAME}" \
    --assume-role-policy-document "${LAMBDA_TRUST_POLICY}" \
    --query "Role.Arn" \
    --output text 2>/dev/null)

if [ $? -ne 0 ]; then
    # Role likely exists, try to get its ARN
    LAMBDA_ROLE_ARN=$(aws iam get-role \
        --role-name "${LAMBDA_ROLE_NAME}" \
        --query "Role.Arn" \
        --output text)
    echo "   IAM Role '${LAMBDA_ROLE_NAME}' already exists: ${LAMBDA_ROLE_ARN}"
else
    echo "   Created IAM Role: ${LAMBDA_ROLE_ARN}"
fi

# Attach the AWS managed policy for basic Lambda execution (CloudWatch Logs access).
aws iam attach-role-policy \
    --role-name "${LAMBDA_ROLE_NAME}" \
    --policy-arn "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"

# Define an inline policy for EC2 start/stop actions. This is attached directly to the role.
EC2_POLICY_DOCUMENT='{"Version": "2012-10-17","Statement": [{"Effect": "Allow","Action": ["ec2:StartInstances","ec2:StopInstances","ec2:DescribeInstances"],"Resource": "*"}]}'

# Put the inline policy onto the role. This will create or update the policy.
aws iam put-role-policy \
    --role-name "${LAMBDA_ROLE_NAME}" \
    --policy-name "${EC2_POLICY_NAME}" \
    --policy-document "${EC2_POLICY_DOCUMENT}"

echo "   Attached EC2 Start/Stop policy to role '${LAMBDA_ROLE_NAME}'."

# Give IAM time to propagate policies (important for newly created roles/policies)
echo "   Waiting for IAM policy propagation..."
sleep 10

# ==========================================================================
# STEP 2: Create/Update the Lambda Function.
# This function will perform the actual EC2 start/stop operations.
# ==========================================================================
echo "
>>> Step 2: Creating/Updating Lambda Function '${LAMBDA_FUNCTION_NAME}'..."

# Check if the Lambda function already exists
FUNCTION_EXISTS=$(aws lambda get-function \
    --function-name "${LAMBDA_FUNCTION_NAME}" \
    --region "${AWS_REGION}" \
    --query "Configuration.FunctionArn" \
    --output text 2>/dev/null)

if [ -n "${FUNCTION_EXISTS}" ]; then
    echo "   Lambda function '${LAMBDA_FUNCTION_NAME}' already exists. Updating code and configuration..."
    # Update existing function code
    aws lambda update-function-code \
        --function-name "${LAMBDA_FUNCTION_NAME}" \
        --zip-file fileb://${zip_file} \
        --region "${AWS_REGION}"
    # Update existing function configuration
    aws lambda update-function-configuration \
        --function-name "${LAMBDA_FUNCTION_NAME}" \
        --runtime "python3.9" \
        --handler "index.lambda_handler" \
        --role "${LAMBDA_ROLE_ARN}" \
        --environment "Variables={INSTANCE_TAG_KEY=${INSTANCE_TAG_KEY},INSTANCE_TAG_VALUE=${INSTANCE_TAG_VALUE},AWS_REGION=${AWS_REGION}}" \
        --timeout 30 \
        --memory-size 128 \
        --region "${AWS_REGION}"
    echo "   Updated Lambda function: ${LAMBDA_FUNCTION_NAME}"
else
    echo "   Lambda function '${LAMBDA_FUNCTION_NAME}' does not exist. Creating..."
    # Create new function
    aws lambda create-function \
        --function-name "${LAMBDA_FUNCTION_NAME}" \
        --runtime "python3.9" \
        --handler "index.lambda_handler" \
        --role "${LAMBDA_ROLE_ARN}" \
        --zip-file fileb://${zip_file} \
        --environment "Variables={INSTANCE_TAG_KEY=${INSTANCE_TAG_KEY},INSTANCE_TAG_VALUE=${INSTANCE_TAG_VALUE},AWS_REGION=${AWS_REGION}}" \
        --timeout 30 \
        --memory-size 128 \
        --region "${AWS_REGION}"
    echo "   Created Lambda function: ${LAMBDA_FUNCTION_NAME}"
fi

# Get the ARN of the Lambda function for EventBridge targets
LAMBDA_FUNCTION_ARN=$(aws lambda get-function \
    --function-name "${LAMBDA_FUNCTION_NAME}" \
    --region "${AWS_REGION}" \
    --query "Configuration.FunctionArn" \
    --output text)

# ==========================================================================
# STEP 3: Create/Update EventBridge Rules to trigger the Lambda function.
# Two rules are created: one for stopping and one for starting instances.
# ==========================================================================
echo "
>>> Step 3: Creating/Updating EventBridge Rules..."

# --- Stop Rule ---
echo "   Configuring Stop Rule: ${STOP_RULE_NAME}"
aws events put-rule \
    --name "${STOP_RULE_NAME}" \
    --schedule-expression "${STOP_CRON_EXPRESSION}" \
    --state "ENABLED" \
    --description "Stop EC2 instances with tag ${INSTANCE_TAG_KEY}=${INSTANCE_TAG_VALUE}" \
    --region "${AWS_REGION}"

# Add Lambda function as a target for the Stop Rule
aws events put-targets \
    --rule "${STOP_RULE_NAME}" \
    --targets "Id=1,Arn=${LAMBDA_FUNCTION_ARN},Input={"action":"stop"}" \
    --region "${AWS_REGION}"

# Grant permission for EventBridge to invoke the Lambda function for the Stop Rule
STOP_RULE_ARN=$(aws events describe-rule --name "${STOP_RULE_NAME}" --region "${AWS_REGION}" --query "Arn" --output text)
aws lambda add-permission \
    --function-name "${LAMBDA_FUNCTION_NAME}" \
    --statement-id "AllowExecutionFromEventBridge-${STOP_RULE_NAME}" \
    --action "lambda:InvokeFunction" \
    --principal "events.amazonaws.com" \
    --source-arn "${STOP_RULE_ARN}" \
    --region "${AWS_REGION}"
echo "   Stop Rule '${STOP_RULE_NAME}' configured."

# --- Start Rule ---
echo "
   Configuring Start Rule: ${START_RULE_NAME}"
aws events put-rule \
    --name "${START_RULE_NAME}" \
    --schedule-expression "${START_CRON_EXPRESSION}" \
    --state "ENABLED" \
    --description "Start EC2 instances with tag ${INSTANCE_TAG_KEY}=${INSTANCE_TAG_VALUE}" \
    --region "${AWS_REGION}"

# Add Lambda function as a target for the Start Rule
aws events put-targets \
    --rule "${START_RULE_NAME}" \
    --targets "Id=1,Arn=${LAMBDA_FUNCTION_ARN},Input={"action":"start"}" \
    --region "${AWS_REGION}"

# Grant permission for EventBridge to invoke the Lambda function for the Start Rule
START_RULE_ARN=$(aws events describe-rule --name "${START_RULE_NAME}" --region "${AWS_REGION}" --query "Arn" --output text)
aws lambda add-permission \
    --function-name "${LAMBDA_FUNCTION_NAME}" \
    --statement-id "AllowExecutionFromEventBridge-${START_RULE_NAME}" \
    --action "lambda:InvokeFunction" \
    --principal "events.amazonaws.com" \
    --source-arn "${START_RULE_ARN}" \
    --region "${AWS_REGION}"
echo "   Start Rule '${START_RULE_NAME}' configured."

# Clean up temporary zip file
rm "${zip_file}"
rm /tmp/index.py

echo "
=== EC2 power management schedule setup completed. ===
"