AWSTemplateFormatVersion: '2010-09-09'
Description: 'AWS Learning Budget Keeper Stack - Automated resource management and cost controls'

Parameters:
  IAMGroup:
    Type: String
    Default: "developers"
    Description: |
      Choose the name of your IAM group where resource restriction policies will be applied.

    ConstraintDescription: Must be the name of an existing IAM group in your AWS account
    AllowedPattern: '^[a-zA-Z0-9+=,.@_-]+$'
    MaxLength: 128
    MinLength: 1

  AutoShutdownHours:
    Type: String
    Default: "10,17,23"
    Description: |
      Times of day (in UTC) when instance shutdown automation should run, comma-separated (1-5 times allowed).
      Examples: "9,21" (9 AM and 9 PM), "6,14,22" (every 8 hours), "8,12,16,20" (4 times daily).
      
      NOTE: All times are in UTC (24-hour format, 0-23).
      To convert from your local time to UTC, use: https://www.timeanddate.com/worldclock/converter.html

    AllowedPattern: '^([0-9]|1[0-9]|2[0-3])(,([0-9]|1[0-9]|2[0-3])){0,4}$'
    ConstraintDescription: Must be a comma-separated list of valid UTC hours (0-23) with 1-5 entries

  TargetRegions:
    Type: String
    Default: "*"
    Description: |
      AWS regions where resource governance should be applied. 
      Examples: "*" (all regions), "us-east-1,us-west-2", "eu-west-1,eu-central-1,ap-southeast-1"
      
    AllowedPattern: '^(\*|([a-z0-9-]+)(,[a-z0-9-]+)*)$'
    ConstraintDescription: Must be "*" for all regions or comma-separated list of valid AWS region names

Resources:
  # IAM Policies for Existing Restricted Group
  DenyDynamoDBReservedCapacityPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: DenyDynamoDBReservedCapacityPurchases
      Groups:
        - !Ref IAMGroup
      PolicyDocument: | 
        {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Sid": "DenyReservedCapacityPurchases",
              "Effect": "Deny",
              "Action": "dynamodb:PurchaseReservedCapacityOfferings",
              "Resource": "arn:aws:dynamodb:*:*:*"
            }
          ]
        }

  LimitEC2Policy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: LimitEC2
      Groups:
        - !Ref IAMGroup
      PolicyDocument: |
        {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Sid": "DenyLargeInstanceTypes",
              "Effect": "Deny",
              "Action": [
                "ec2:RunInstances",
                "ec2:ModifyInstanceAttribute"
              ],
              "Resource": [
                "arn:aws:ec2:*:*:instance/*"
              ],
              "Condition": {
                "ForAnyValue:StringNotLike": {
                  "ec2:InstanceType": [
                    "t2.nano",
                    "t2.micro",
                    "t2.small",
                    "t2.medium",
                    "t3.nano",
                    "t3.micro",
                    "t3.small",
                    "t3.medium",
                    "t4g.nano",
                    "t4g.micro",
                    "t4g.small",
                    "t4g.medium"
                  ]
                }
              }
            },
            {
              "Sid": "DenyReservedInstances",
              "Effect": "Deny",
              "Action": [
                "ec2:ModifyReservedInstances",
                "ec2:PurchaseReservedInstancesOffering"
              ],
              "Resource": "*"
            },
            {
              "Sid": "LimitInstanceVolumeType",
              "Effect": "Deny",
              "Action": "ec2:*",
              "Resource": "*",
              "Condition": {
                "ForAnyValue:StringNotLike": {
                  "ec2:VolumeType": [
                    "gp2",
                    "gp3",
                    "standard"
                  ]
                }
              }
            },
            {
              "Sid": "LimitVolumeSize",
              "Effect": "Deny",
              "Action": "ec2:*",
              "Resource": "*",
              "Condition": {
                "NumericGreaterThan": {
                  "ec2:VolumeSize": "30"
                }
              }
            },
            {
              "Sid": "LimitVolumeModificationSize",
              "Effect": "Deny",
              "Action": "ec2:*",
              "Resource": "*",
              "Condition": {
                "NumericGreaterThan": {
                  "ec2:TargetVolumeSize": "30"
                }
              }
            }
          ]
        }

  # IAM Role for Lambda Function
  LearningBudgetKeeperLambdaRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: LearningBudgetKeeperLambdaRole
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: BudgetProtectionPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - ec2:DescribeInstances
                  - ec2:DescribeInstanceAttribute
                  - ec2:DescribeTags
                  - ec2:StopInstances
                  - ec2:StartInstances
                  - ec2:DescribeInstanceStatus
                  - ec2:DescribeRegions
                  - autoscaling:DescribeAutoScalingGroups
                  - autoscaling:UpdateAutoScalingGroup
                  - autoscaling:DescribeScalingActivities
                  - autoscaling:DescribeTags
                Resource: '*'

  # EventBridge Rule for configurable execution schedule
  LearningBudgetKeeperScheduleRule:
    Type: AWS::Events::Rule
    Properties:
      Name: learning-budget-keeper-schedule
      Description: Trigger budget protection Lambda at specified UTC times
      ScheduleExpression: !Sub 'cron(0 ${AutoShutdownHours} * * ? *)'
      State: ENABLED
      Targets:
        - Arn: !GetAtt LearningBudgetKeeperLambda.Arn
          Id: LearningBudgetKeeperLambdaTarget

  # Permission for EventBridge to invoke Lambda
  LambdaInvokePermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !Ref LearningBudgetKeeperLambda
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt LearningBudgetKeeperScheduleRule.Arn

  # CloudWatch Log Group for Lambda
  LearningBudgetKeeperLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub '/aws/lambda/${LearningBudgetKeeperLambda}'
      RetentionInDays: 30

  # Lambda Function for Budget Protection
  LearningBudgetKeeperLambda:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: aws-learning-budget-keeper-function
      Runtime: python3.12
      Handler: index.lambda_handler
      Role: !GetAtt LearningBudgetKeeperLambdaRole.Arn
      Timeout: 300
      Environment:
        Variables:
          TARGET_REGIONS: !Ref TargetRegions
      Code:
        ZipFile: |
          """
          AWS Learning Budget Keeper Lambda Function
          
          This function automatically manages AWS resources to control costs and enforce governance:
          - Stops running EC2 instances (except protected ones)
          - Scales down Auto Scaling Groups to zero
          
          """
          
          import boto3
          import json
          import os
          from datetime import datetime, timezone
          from typing import List, Dict, Any
          
          
          def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
              print(f"Learning Budget Keeper Lambda triggered at {datetime.now(timezone.utc)}")
          
              # Get target regions from environment variable
              target_regions_env = os.environ.get('TARGET_REGIONS', '*')
              print(f"Target regions configuration: {target_regions_env}")
              
              # Always get available regions first for validation
              try:
                  ec2_client = boto3.client('ec2')
                  regions_response = ec2_client.describe_regions(
                      Filters=[
                          {'Name': 'opt-in-status', 'Values': ['opt-in-not-required', 'opted-in']}
                      ]
                  )
                  available_regions = [region['RegionName'] for region in regions_response['Regions']]
                  print(f"Available enabled regions: {available_regions}")
              except Exception as e:
                  error_msg = f"Failed to discover available regions: {str(e)}"
                  print(error_msg)
                  return {
                      'statusCode': 500,
                      'body': json.dumps({
                          'error': error_msg,
                          'message': 'Resource governance failed - could not discover available regions'
                      })
                  }
              
              if target_regions_env == '*':
                  # Use all available enabled regions
                  allowed_regions = available_regions
                  print(f"Using all available regions: {allowed_regions}")
              else:
                  # Parse and validate specified regions
                  specified_regions = [region.strip() for region in target_regions_env.split(',') if region.strip()]
                  
                  # Validate that all specified regions are available
                  invalid_regions = [region for region in specified_regions if region not in available_regions]
                  if invalid_regions:
                      error_msg = f"Invalid or unavailable regions specified: {invalid_regions}. Available regions: {available_regions}"
                      print(error_msg)
                      return {
                          'statusCode': 400,
                          'body': json.dumps({
                              'error': error_msg,
                              'message': 'Resource governance failed - invalid regions specified'
                          })
                      }
                  
                  allowed_regions = specified_regions
                  print(f"Using validated specified regions: {allowed_regions}")
          
              print(f"Processing resource governance across regions: {allowed_regions}")
          
              # Aggregate results across all regions
              aggregate_results = {
                  'stopped_instances': [],
                  'scaled_down_asgs': [],
                  'errors': [],
                  'regions_processed': []
              }
          
              # Process each allowed region
              for region in allowed_regions:
                  try:
                      print(f"Processing region: {region}")
                      region_results = process_region(region)
          
                      # Aggregate results
                      aggregate_results['stopped_instances'].extend(region_results['stopped_instances'])
                      aggregate_results['scaled_down_asgs'].extend(region_results['scaled_down_asgs'])
                      aggregate_results['errors'].extend(region_results['errors'])
                      aggregate_results['regions_processed'].append(region)
          
                  except Exception as e:
                      error_msg = f"Error processing region {region}: {str(e)}"
                      print(error_msg)
                      aggregate_results['errors'].append(error_msg)
          
              return {
                  'statusCode': 200,
                  'body': json.dumps(aggregate_results)
              }
          
          
          def process_region(region: str) -> Dict[str, List[str]]:
              """
              Process resource governance for a specific region.
          
              Args:
                  region: AWS region name
          
              Returns:
                  Dict containing regional results
              """
              print(f"Initializing AWS clients for region: {region}")
          
              # Initialize region-specific AWS clients
              ec2 = boto3.client('ec2', region_name=region)
              autoscaling = boto3.client('autoscaling', region_name=region)
          
              region_results = {
                  'stopped_instances': [],
                  'scaled_down_asgs': [],
                  'errors': []
              }
          
              try:
                  # 1. Stop running EC2 instances (except those with specific tags)
                  stopped_instances = stop_idle_instances(ec2, region)
                  region_results['stopped_instances'] = stopped_instances
          
                  # 2. Scale down Auto Scaling Groups to zero
                  scaled_down_asgs = scale_down_asgs(autoscaling, region)
                  region_results['scaled_down_asgs'] = scaled_down_asgs
          
                  print(f"Region {region} processing completed - Instances: {len(stopped_instances)}, ASGs: {len(scaled_down_asgs)}")
          
              except Exception as e:
                  error_msg = f"Error processing region {region}: {str(e)}"
                  print(error_msg)
                  region_results['errors'].append(error_msg)
          
              return region_results
          
          
          def is_instance_stop_protected(ec2, instance_id: str) -> bool:
              """
              Check if an EC2 instance has stop protection enabled via API or custom tags.
          
              Args:
                  ec2: Boto3 EC2 client
                  instance_id: EC2 instance ID
          
              Returns:
                  True if instance has stop protection, False otherwise
              """
              try:
                  # Check API stop protection
                  response = ec2.describe_instance_attribute(
                      InstanceId=instance_id,
                      Attribute='disableApiStop'
                  )
                  if response.get('DisableApiStop', {}).get('Value', False):
                      return True
          
                  # Check custom ResourceGovernance tag
                  tags_response = ec2.describe_tags(
                      Filters=[
                          {'Name': 'resource-id', 'Values': [instance_id]},
                          {'Name': 'key', 'Values': ['ResourceGovernance']}
                      ]
                  )
          
                  for tag in tags_response.get('Tags', []):
                      if tag.get('Value', '').lower() == 'keep':
                          return True
          
                  return False
          
              except Exception as e:
                  print(f"Error checking stop protection for instance {instance_id}: {str(e)}")
                  # If we can't determine protection status, assume it's protected to be safe
                  return True
          
          
          def is_asg_protected(autoscaling, asg_name: str) -> bool:
              """
              Check if an Auto Scaling Group has protection via custom tags.
          
              Args:
                  autoscaling: Boto3 Auto Scaling client
                  asg_name: Auto Scaling Group name
          
              Returns:
                  True if ASG has protection, False otherwise
              """
              try:
                  response = autoscaling.describe_tags(
                      Filters=[
                          {'Name': 'auto-scaling-group', 'Values': [asg_name]},
                          {'Name': 'key', 'Values': ['ResourceGovernance']}
                      ]
                  )
          
                  for tag in response.get('Tags', []):
                      if tag.get('Value', '').lower() == 'keep':
                          return True
          
                  return False
          
              except Exception as e:
                  print(f"Error checking protection for ASG {asg_name}: {str(e)}")
                  # If we can't determine protection status, assume it's protected to be safe
                  return True
          
          
          def stop_idle_instances(ec2, region: str) -> List[str]:
              """
              Stop all running EC2 instances except those with stop protection enabled.
          
              Args:
                  ec2: Boto3 EC2 client
                  region: AWS region name
          
              Returns:
                  List of stopped instance IDs
              """
              stopped_instances = []
          
              try:
                  # Get all running instances
                  response = ec2.describe_instances(
                      Filters=[
                          {'Name': 'instance-state-name', 'Values': ['running']}
                      ]
                  )
          
                  for reservation in response['Reservations']:
                      for instance in reservation['Instances']:
                          instance_id = instance['InstanceId']
                          instance_type = instance.get('InstanceType', 'unknown')
          
                          # Check if instance has stop protection
                          if is_instance_stop_protected(ec2, instance_id):
                              print(f"Skipping protected instance: {instance_id} (type: {instance_type}) in region {region}")
                              continue
          
                          try:
                              ec2.stop_instances(InstanceIds=[instance_id])
                              stopped_instances.append(f"{instance_id} ({region})")
                              print(f"Stopped instance: {instance_id} (type: {instance_type}) in region {region}")
                          except Exception as e:
                              error_msg = f"Failed to stop instance {instance_id} in {region}: {str(e)}"
                              print(error_msg)
          
              except Exception as e:
                  print(f"Error stopping instances in {region}: {str(e)}")
          
              return stopped_instances
          
          
          def scale_down_asgs(autoscaling, region: str) -> List[str]:
              """
              Scale down all Auto Scaling Groups to zero capacity except protected ones.
          
              Args:
                  autoscaling: Boto3 Auto Scaling client
                  region: AWS region name
          
              Returns:
                  List of scaled down ASG names
              """
              scaled_down_asgs = []
          
              try:
                  # Get all Auto Scaling Groups
                  paginator = autoscaling.get_paginator('describe_auto_scaling_groups')
          
                  for page in paginator.paginate():
                      for asg in page['AutoScalingGroups']:
                          asg_name = asg['AutoScalingGroupName']
                          current_capacity = asg['DesiredCapacity']
          
                          # Check if ASG is protected
                          if is_asg_protected(autoscaling, asg_name):
                              print(f"Skipping protected ASG: {asg_name} (capacity: {current_capacity}) in region {region}")
                              continue
          
                          if current_capacity > 0:
                              try:
                                  autoscaling.update_auto_scaling_group(
                                      AutoScalingGroupName=asg_name,
                                      MinSize=0,
                                      MaxSize=asg['MaxSize'],
                                      DesiredCapacity=0
                                  )
                                  scaled_down_asgs.append(f"{asg_name} ({region})")
                                  print(f"Scaled down ASG: {asg_name} (from {current_capacity} to 0) in region {region}")
                              except Exception as e:
                                  error_msg = f"Failed to scale down ASG {asg_name} in {region}: {str(e)}"
                                  print(error_msg)
                          else:
                              print(f"ASG already at zero capacity: {asg_name} in region {region}")
          
              except Exception as e:
                  print(f"Error scaling down ASGs in {region}: {str(e)}")
          
              return scaled_down_asgs