AWSTemplateFormatVersion: "2010-09-09"
Description: AWS Backup Restore Testing Plan Configuration

Parameters:
  RDSSubnetGroupParam:
    Description: "Optional: Provide an RDS Subnet Group Name. Leave blank to auto-select one."
    Type: String
    Default: ""

Conditions:
  HasCustomSubnetGroup: !Not [ !Equals [ !Ref RDSSubnetGroupParam, "" ] ]

Resources:

  # Step 1 IAM Roles creation
  # Role for AWS Backup operations used by restore testing selections.
  BackupServiceRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: BackupServiceRole
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - backup.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSBackupServiceRolePolicyForBackup
        - arn:aws:iam::aws:policy/service-role/AWSBackupServiceRolePolicyForRestores
        - arn:aws:iam::aws:policy/AWSBackupServiceRolePolicyForS3Backup
        - arn:aws:iam::aws:policy/AWSBackupServiceRolePolicyForS3Restore

  # Role for the RDS Subnet Finder custom resource Lambda.
  LambdaSubnetFinderRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: LambdaSubnetFinderRole
      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: RDSSubnetFinderPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - rds:DescribeDBSubnetGroups
                Resource: "*"

  # Role for the Restore Coordinator Lambda.
  LambdaCoordinatorRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: LambdaCoordinatorRole
      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: LambdaInvokePolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - lambda:InvokeFunction
                Resource: "*"  

  # Role for the S3 Restore Validation Lambda.
  LambdaS3RestoreRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: LambdaS3RestoreRole
      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: S3RestoreValidationPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - s3:ListBucket
                  - s3:ListBucketVersions
                  - s3:ListAllMyBuckets
                  - s3:GetObject
                Resource: "*"
              - Effect: Allow
                Action:
                  - backup:PutRestoreValidationResult
                Resource: "*"

  # Role for the RDS Restore Validation Lambda.
  LambdaRDSRestoreRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: LambdaRDSRestoreRole
      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: RDSRestoreValidationPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - rds:DescribeDBInstances
                Resource: "*"
              - Effect: Allow
                Action:
                  - backup:PutRestoreValidationResult
                Resource: "*"

  # Step 2: RDS Lambda Subnet Finder
  RDSSubnetGroupFinderLambda:
    Type: AWS::Lambda::Function
    DependsOn: LambdaSubnetFinderRole
    Properties:
      FunctionName: RDSSubnetGroupFinder
      Handler: index.handler
      Role: !GetAtt LambdaSubnetFinderRole.Arn
      Code:
        ZipFile: |
          import json
          import boto3
          import logging
          import urllib3

          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          http = urllib3.PoolManager()

          def send_response(event, context, response_status, response_data, physical_resource_id=None, no_echo=False):
              response_url = event['ResponseURL']
              response_body = {
                  'Status': response_status,
                  'Reason': 'See details in CloudWatch Log Stream: ' + context.log_stream_name,
                  'PhysicalResourceId': physical_resource_id or context.log_stream_name,
                  'StackId': event['StackId'],
                  'RequestId': event['RequestId'],
                  'LogicalResourceId': event['LogicalResourceId'],
                  'Data': response_data
              }
              json_response_body = json.dumps(response_body)
              headers = {
                  'content-type': '',
                  'content-length': str(len(json_response_body))
              }
              try:
                  response = http.request('PUT', response_url, body=json_response_body, headers=headers)
                  logger.info("CloudFormation response status: " + str(response.status))
              except Exception as e:
                  logger.error("send_response failed: " + str(e))

          def handler(event, context):
              logger.info("Received event: " + json.dumps(event))
              try:
                  client = boto3.client('rds')
                  response = client.describe_db_subnet_groups()
                  subnet_groups = response.get('DBSubnetGroups', [])
                  if subnet_groups:
                      import random
                      chosen = random.choice(subnet_groups)['DBSubnetGroupName']
                  else:
                      chosen = ""
                  logger.info("Auto-selected subnet group: " + chosen)
                  send_response(event, context, "SUCCESS", {"SubnetGroupName": chosen})
              except Exception as e:
                  logger.error("Exception occurred: " + str(e))
                  send_response(event, context, "FAILED", {"Message": str(e)})
      Runtime: python3.12
      Timeout: 60
      MemorySize: 128

  # Custom Resource to invoke the RDS Subnet Finder Lambda
  RDSSubnetGroupFinder:
    Type: AWS::CloudFormation::CustomResource
    DependsOn: RDSSubnetGroupFinderLambda
    Properties:
      ServiceToken: !GetAtt RDSSubnetGroupFinderLambda.Arn

  # Step 3: Create Restore Testing Plan with a daily schedule at 10PM (UTC)
  BackupRestoreTestingPlan:
    Type: AWS::Backup::RestoreTestingPlan
    DependsOn: RDSSubnetGroupFinder
    Properties:
      ScheduleExpression: cron(0 22 ? * * *)
      RecoveryPointSelection:
        RecoveryPointTypes:
          - SNAPSHOT
        IncludeVaults:
          - '*'
        Algorithm: LATEST_WITHIN_WINDOW
        SelectionWindowDays: 7
      StartWindowHours: 4
      RestoreTestingPlanName: DailyRestorePlan

  # Step 4: S3 Restore Testing Selection
  BackupRestoreTestingSelectionS3:
    Type: AWS::Backup::RestoreTestingSelection
    DependsOn: BackupRestoreTestingPlan
    Properties:
      ProtectedResourceType: S3
      RestoreTestingSelectionName: RestoreTestingSelectionS3
      RestoreTestingPlanName: DailyRestorePlan
      IamRoleArn: !GetAtt BackupServiceRole.Arn
      ValidationWindowHours: 4

  # Step 5: RDS Restore Testing Selection using conditional override
  BackupRestoreTestingSelectionRDS:
    Type: AWS::Backup::RestoreTestingSelection
    DependsOn: BackupRestoreTestingPlan
    Properties:
      ProtectedResourceType: RDS
      RestoreTestingSelectionName: RestoreTestingSelectionRDS
      RestoreTestingPlanName: DailyRestorePlan
      IamRoleArn: !GetAtt BackupServiceRole.Arn
      ValidationWindowHours: 4
      RestoreMetadataOverrides: !If
        - HasCustomSubnetGroup
        - { dbSubnetGroupName: !Ref RDSSubnetGroupParam }
        - { dbSubnetGroupName: !GetAtt RDSSubnetGroupFinder.SubnetGroupName }

  # Step 6: Restore Coordinator Lambda (invokes S3 or RDS validation)
  RestoreCoordinatorLambda:
    Type: AWS::Lambda::Function
    DependsOn: LambdaCoordinatorRole
    Properties:
      FunctionName: RestoreValidationCoordinator
      Handler: index.handler
      Role: !GetAtt LambdaCoordinatorRole.Arn
      Code:
        ZipFile: |
          import json
          import boto3
          import logging
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          lambda_client = boto3.client('lambda')
          def handler(event, context):
              logger.info("Handling event: %s", json.dumps(event))
              resource_type = event.get('detail', {}).get('resourceType', '')
              function_name = None
              try:
                  if resource_type == "RDS":
                      function_name = "RDSRestoreValidation"
                      logger.info("Resource is an RDS instance. Invoking Lambda function: %s", function_name)
                  elif resource_type == "S3":
                      function_name = "S3RestoreValidation"
                      logger.info("Resource is an S3 bucket. Invoking Lambda function: %s", function_name)
                  else:
                      raise ValueError(f"Unsupported resource type: {resource_type}")
                  response = lambda_client.invoke(
                      FunctionName=function_name,
                      Payload=json.dumps(event),
                      InvocationType="RequestResponse"
                  )
                  logger.info("Lambda invoke response: %s", response)
              except Exception as e:
                  logger.error("Error during Lambda invocation: %s", str(e))
                  raise e
              logger.info("Finished processing event for resource type: %s", resource_type)
      Runtime: python3.12
      Timeout: 60
      MemorySize: 128

  # Step 7: Create EventBridge Rule to trigger the Restore Coordinator Lambda
  BackupRestoreTestingEventRule:
    Type: AWS::Events::Rule
    DependsOn: 
      - BackupRestoreTestingPlan
      - RestoreCoordinatorLambda
    Properties:
      Name: Backup_restore_testing
      EventPattern:
        {"source":["aws.backup"],"detail-type":["Restore Job State Change"],"detail":{"status":["COMPLETED"],"restoreTestingPlanArn":[{"prefix":!GetAtt BackupRestoreTestingPlan.RestoreTestingPlanArn}]}}
      State: ENABLED
      EventBusName: default
      Targets:
        - Id: LambdaTarget
          Arn: !GetAtt RestoreCoordinatorLambda.Arn

  # Step 8: Lambda Permission for EventBridge
  LambdaInvokePermission:
    Type: AWS::Lambda::Permission
    DependsOn: BackupRestoreTestingEventRule
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !GetAtt RestoreCoordinatorLambda.Arn
      Principal: events.amazonaws.com
      SourceArn: !GetAtt BackupRestoreTestingEventRule.Arn

  # Step 9: S3 Restore Validation Lambda
  S3ValidationLambda:
    Type: AWS::Lambda::Function
    DependsOn: LambdaS3RestoreRole
    Properties:
      FunctionName: S3RestoreValidation
      Handler: index.handler
      Role: !GetAtt LambdaS3RestoreRole.Arn
      Code:
        ZipFile: |
          import json
          import boto3
          import logging
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          s3_client = boto3.client('s3')
          backup_client = boto3.client('backup')
          def handler(event, context):
              logger.info("Handling event: %s", json.dumps(event))
              restore_job_id = event.get('detail', {}).get('restoreJobId', '')
              resource_type = event.get('detail', {}).get('resourceType', '')
              created_resource_arn = event.get('detail', {}).get('createdResourceArn', '')
              validation_status = "SUCCESSFUL"
              validation_status_message = "Restore validation completed successfully"
              try:
                  if resource_type == "S3":
                      bucket_name = get_bucket_name_from_arn(created_resource_arn)
                      response = s3_client.list_objects_v2(Bucket=bucket_name)
                      object_count = response.get('KeyCount', 0)
                      if object_count > 1:
                          logger.info(f"Bucket {bucket_name} contains more than 1 object. Validation successful.")
                      else:
                          logger.info(f"Bucket {bucket_name} contains 1 or fewer objects. Validation failed.")
                          validation_status = "FAILED"
                          validation_status_message = f"Bucket {bucket_name} contains only {object_count} object(s)."
                  else:
                      validation_status = "FAILED"
                      validation_status_message = f"Unsupported resource type: {resource_type}"
                  backup_client.put_restore_validation_result(
                      RestoreJobId=restore_job_id,
                      ValidationStatus=validation_status,
                      ValidationStatusMessage=validation_status_message
                  )
                  logger.info("Restore validation result sent successfully")
              except Exception as e:
                  logger.error("Error during restore validation: %s", str(e))
                  validation_status = "FAILED"
                  validation_status_message = f"Restore validation encountered an error: {str(e)}"
                  backup_client.put_restore_validation_result(
                      RestoreJobId=restore_job_id,
                      ValidationStatus=validation_status,
                      ValidationStatusMessage=validation_status_message
                  )
              logger.info("Finished processing restore validation for job ID: %s", restore_job_id)
          def get_bucket_name_from_arn(arn):
              arn_parts = arn.split(":")
              resource_parts = arn_parts[-1].split("/")
              return resource_parts[-1]
      Runtime: python3.12
      Timeout: 60
      MemorySize: 128

  # Step 10: RDS Restore Validation Lambda
  RDSValidationLambda:
    Type: AWS::Lambda::Function
    DependsOn: LambdaRDSRestoreRole
    Properties:
      FunctionName: RDSRestoreValidation
      Handler: index.handler
      Role: !GetAtt LambdaRDSRestoreRole.Arn
      Code:
        ZipFile: |
          import json
          import boto3
          import logging
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          rds_client = boto3.client('rds')
          backup_client = boto3.client('backup')
          def handler(event, context):
              logger.info("Handling event: %s", json.dumps(event))
              restore_job_id = event.get('detail', {}).get('restoreJobId', '')
              resource_type = event.get('detail', {}).get('resourceType', '')
              created_resource_arn = event.get('detail', {}).get('createdResourceArn', '')
              validation_status = "SUCCESSFUL"
              validation_status_message = "Restore validation completed successfully"
              try:
                  if resource_type == "RDS":
                      instance_id = get_instance_id_from_arn(created_resource_arn)
                      describe_response = rds_client.describe_db_instances(DBInstanceIdentifier=instance_id)
                      db_instance = describe_response['DBInstances'][0]
                      db_status = db_instance['DBInstanceStatus']
                      logger.info(f'RDS instance status: {db_status}')
                      if db_status != "available":
                          validation_status = "FAILED"
                          validation_status_message = f"RDS instance {instance_id} is not available, current status: {db_status}"
                  backup_client.put_restore_validation_result(
                      RestoreJobId=restore_job_id,
                      ValidationStatus=validation_status,
                      ValidationStatusMessage=validation_status_message
                  )
                  logger.info("Restore validation result sent successfully")
              except Exception as e:
                  logger.error("Error during restore validation: %s", str(e))
                  validation_status = "FAILED"
                  validation_status_message = f"Restore validation encountered an error: {str(e)}"
                  backup_client.put_restore_validation_result(
                      RestoreJobId=restore_job_id,
                      ValidationStatus=validation_status,
                      ValidationStatusMessage=validation_status_message
                  )
              logger.info("Finished processing restore validation for job ID: %s", restore_job_id)
          def get_instance_id_from_arn(arn):
              arn_parts = arn.split(":")
              resource_parts = arn_parts[-1].split("/")
              return resource_parts[-1]
      Runtime: python3.12
      Timeout: 60
      MemorySize: 128