AWSTemplateFormatVersion: '2010-09-09'
Description: 'Deploys resources for simulating partial failures using ALB, Lambda, and AWS FIS'

Parameters:
  LoadBalancerName:
    Type: String
    Description: 'Name of the Load Balancer'
  ListenerRuleArn:
    Type: String
    Description: 'ARN of the ALB listener rule to modify during the failure simulation'
  TestDurationInSeconds:
    Type: Number
    Description: 'Duration (in seconds) for how long the failure should persist'
Resources:

  FISExperimentLogGroup:
    Type: 'AWS::Logs::LogGroup'
    Metadata:
      cfn_nag:
        rules_to_suppress:
          - id: W84
            reason: "It is recommended that CloudWatchLogs LogGroup should specify a KMS Key Id to encrypt the log data, but because this is a sample template we can ignore this"
    Properties:
      LogGroupName: !Sub '${AWS::StackName}-Log'  # You can customize this name
      RetentionInDays: 7             # Retention period in days
      
  ErrorSimulationFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: 
                - lambda.amazonaws.com
                - elasticloadbalancing.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  ErrorSimulationFunction:
    Type: AWS::Lambda::Function
    Metadata:
      cfn_nag:
        rules_to_suppress:
          - id: W92
            reason: "This is a sample Lambda where we don't need to reserved execution concurrency"
          - id: W89
            reason: "This lambda is used to execute call AWS services, so we don't need to run inside VPC unless required by other security policies"
    Properties:
      FunctionName: !Sub ${AWS::StackName}-error-simulation
      Handler: index.lambda_handler
      Role: !GetAtt ErrorSimulationFunctionRole.Arn
      Code:
        ZipFile: |
          def lambda_handler(event, context):
              return {
                  "statusCode": 500,
                  "statusDescription": "500 Internal Server Error", 
                  "isBase64Encoded": False,
                  "headers": {
                      "Content-Type": "text/html"
                  },
                  "body": "💥 Service Unavailable - Simulated Error Response\n"
              }
      Runtime: python3.9
      Timeout: 30

  TargetGroupFailureSimulationDocument:
    Type: AWS::SSM::Document
    Properties:
      DocumentType: Automation
      Name: TargetGroup-Failure-Simulation-FIS
      Content:
        schemaVersion: '0.3'
        description: |
          ### Document Name -TargetGroup-Failure-Simulation-FIS

          ## What does this document do? 

          This document creates a new target group that points to a Lambda function configured to return specific error codes. 
          It then modifies the specified listener rule in the provided Application Load Balancer (ALB) to route a portion of traffic to this new target group using a configurable weight.
          After waiting for the specified test duration, the target group is removed from the ALB listener. If the rule is still active after the maximum duration, 
          it will be forcefully revoked. The automation then completes execution once it confirms the cleanup is successful.

          
          ## Security Risk
          Low: The Lambda function is configured to return a fixed, non-malicious response (e.g., 503) and does not process or store any user data.


          ## Input Parameters
          * LambdaFunctionArn: (Required) The ARN of the Lambda function that simulates failure responses
          * AutomationAssumeRole: (Optional) The ARN of the IAM role that allows AWS Systems Manager Automation to perform actions on your behalf
          * TargetGroupName: (Required) The name of the target group that routes traffic to the Lambda function.
          * TestDurationInMinutes: (Optional) ** Default 1 minute ** – The duration to keep the test rule active before rollback is triggered.

          ## Supports Rollback
          Yes. The target group is detached from the listener and deleted as part of the rollback process.

          ## Cancellation behaviour
          Yes. If the automation is cancelled, the target group is removed from the ALB listener and deleted to restore the original configuration.

          ## Output Parameters
          This document has no outputs.

          ## Minimum Permissions Required
          * `elbv2:CreateTargetGroup`
          * `elbv2:ModifyListener`
          * `lambda:AddPermission`

          ## Additional Permissions for logging
          * logs:CreateLogStream
          * logs:CreateLogGroup
          * logs:PutLogEvents
          * logs:DescribeLogGroups
          * logs:DescribeLogStreams
        assumeRole: '{{ AutomationAssumeRole }}'
        parameters:
          AutomationAssumeRole:
            default: ''
            description: (Optional) The ARN of the role that allows Automation to perform the actions on your behalf.
            type: String
            allowedPattern: ^arn:aws(-cn|-us-gov)?:iam::\d{12}:role\/[\w+=,.@_\/-]+|^$
          RuleArn:
            description: ARN of the Rule to modify
            type: String
          LambdaFunctionName:
            description: Name of the Function To Invoke
            type: String
          DurationSeconds:
            default: '1'
            description: (Optional) Duration in seconds.
            type: String
          FailurePercentage:
            description: From 0 to 100 what is the percentage of errors that should be returned
            type: Integer
        variables:
          IsDefaultRule:
            default: 'False'
            description: Used to detect if ALB Rule is the default rule, because updating this rule follows a different API mechanism (using ModifyListener)
            type: String
          LambdaFunctionArn:
            description: Constructed ARN for a given Lambda Function Name
            type: String
          TargetGroupName:
            description: Temporary TargetGroupName for the duration of the experiment
            type: String
        mainSteps:
          - description: Helper to create a random TargetGroupName
            name: CustomVariablesHelper
            action: aws:executeScript
            nextStep: UpdateTargetGroupName
            isEnd: false
            inputs:
              Script: |-
                import uuid

                def script_handler(events, context):
                  randomSuffix = str(uuid.uuid4())[-12:]
                  targetGroupName = f'fisautomation-{randomSuffix}-tg'
                  
                  return {'TargetGroupName': targetGroupName }
              Runtime: python3.11
              InputPayload:
                LambdaFunctionName: '{{ LambdaFunctionName }}'
                TargetGroupName: '{{ variable:TargetGroupName }}'
              Handler: script_handler
            outputs:
              - Type: String
                Name: TargetGroupName
                Selector: $.Payload.TargetGroupName
          - description: Uses the Random Name created by the script and updates the Document Variable
            name: UpdateTargetGroupName
            action: aws:updateVariable
            nextStep: CreateLambdaARN
            isEnd: false
            inputs:
              Value: '{{ CustomVariablesHelper.TargetGroupName }}'
              Name: variable:TargetGroupName
          - description: Constructs the LambdaARN based on the Function Name passed as parameter
            name: CreateLambdaARN
            action: aws:updateVariable
            nextStep: CreateTargetGroup
            isEnd: false
            inputs:
              Value: arn:{{ global:AWS_PARTITION }}:lambda:{{ global:REGION }}:{{ global:ACCOUNT_ID }}:function:{{ LambdaFunctionName }}
              Name: variable:LambdaFunctionArn
          - description: Create Target Group
            name: CreateTargetGroup
            action: aws:executeAwsApi
            nextStep: AddPermission
            isEnd: false
            onFailure: Abort
            inputs:
              TargetType: lambda
              Service: elbv2
              Api: CreateTargetGroup
              Name: '{{ variable:TargetGroupName }}'
            outputs:
              - Name: TargetGroupArn
                Selector: $.TargetGroups[0].TargetGroupArn
          - name: AddPermission
            action: aws:executeAwsApi
            nextStep: RegisterTargets
            isEnd: false
            inputs:
              FunctionName: '{{ LambdaFunctionName }}'
              Action: lambda:InvokeFunction
              SourceArn: '{{ CreateTargetGroup.TargetGroupArn }}'
              StatementId: SID{{ automation:EXECUTION_ID }}
              Service: lambda
              Api: AddPermission
              Principal: elasticloadbalancing.amazonaws.com
          - description: Register Targets
            name: RegisterTargets
            action: aws:executeAwsApi
            maxAttempts: 3
            nextStep: DescribeRules
            isEnd: false
            onFailure: step:RemovePermission
            inputs:
              TargetGroupArn: '{{ CreateTargetGroup.TargetGroupArn }}'
              Service: elbv2
              Targets:
                - Id: '{{ variable:LambdaFunctionArn }}'
              Api: RegisterTargets
          - name: DescribeRules
            action: aws:executeAwsApi
            nextStep: CalculateRuleSettings
            isEnd: false
            inputs:
              RuleArns:
                - '{{RuleArn}}'
              Service: elbv2
              Api: DescribeRules
            outputs:
              - Type: StringMap
                Name: DescribeRulesOutput
                Selector: $
          - name: CalculateRuleSettings
            action: aws:executeScript
            nextStep: UpdateIsDefaultRuleVariable
            isEnd: false
            onCancel: step:RemovePermission
            onFailure: step:RemovePermission
            inputs:
              Script: |-
                import copy

                def script_handler(event, context):
                    old_config = copy.deepcopy(event)
                    new_config = copy.deepcopy(event)
                    # checks before any changes:

                    failure_pct = event['FailurePercentage']
                    target_group_arn = event['TargetGroupARN']

                    if len(event['Payload']['Rules']) > 1:
                        raise Exception("Too many rules in DescribeRules payload. Aborting")

                    if len(event['Payload']['Rules'][0]['Actions']) > 1:
                        raise Exception("Too many actions in DescribeRules payload. Aborting")

                    if len(event['Payload']['Rules'][0]['Actions'][0]['ForwardConfig']['TargetGroups']) > 1:
                        raise Exception("Too many TargetGroups in DescribeRules payload. Aborting")

                    if event['Payload']['Rules'][0]['Actions'][0]['ForwardConfig']['TargetGroups'][0]['Weight'] != 1:
                        raise Exception("Existing Weight in DescribeRules payload different from default. Aborting")

                    # Modify config

                    del new_config['Payload']['Rules'][0]['Actions'][0]['TargetGroupArn']

                    new_config['Payload']['Rules'][0]['Actions'][0]['ForwardConfig']['TargetGroups'][0]['Weight'] = 100-failure_pct

                    new_config['Payload']['Rules'][0]['Actions'][0]['ForwardConfig']['TargetGroups'].append(
                        {
                            "TargetGroupArn": target_group_arn,
                            "Weight": failure_pct
                        }
                    )
                    
                    # Calculates Listener ARN
                    rule = event['Payload']['Rules'][0]['RuleArn']
                    rule_split = rule.split(":")
                    listener_part = rule_split[5].split("/")
                    listener_part[0]='listener'
                    listener_root_arn = ":".join(rule_split[0:-1])
                    listener_arn = ":".join([listener_root_arn, "/".join(listener_part[0:-1])])


                    return {
                        'OriginalActions': old_config['Payload']['Rules'][0]['Actions'],
                        'Actions': new_config['Payload']['Rules'][0]['Actions'],
                        'IsDefaultRule': str(event['Payload']['Rules'][0]['IsDefault']),
                        'ListenerArn': listener_arn
                    }
              Runtime: python3.11
              InputPayload:
                TargetGroupARN: '{{ CreateTargetGroup.TargetGroupArn }}'
                Payload: '{{ DescribeRules.DescribeRulesOutput }}'
                FailurePercentage: '{{ FailurePercentage }}'
              Handler: script_handler
            outputs:
              - Type: MapList
                Name: Actions
                Selector: $.Payload.Actions
              - Type: MapList
                Name: OriginalActions
                Selector: $.Payload.OriginalActions
              - Type: String
                Name: IsDefaultRule
                Selector: $.Payload.IsDefaultRule
              - Type: String
                Name: ListenerArn
                Selector: $.Payload.ListenerArn
          - name: UpdateIsDefaultRuleVariable
            action: aws:updateVariable
            nextStep: ChooseListenerOrRuleAddTG
            isEnd: false
            inputs:
              Value: '{{ CalculateRuleSettings.IsDefaultRule }}'
              Name: variable:IsDefaultRule
          - name: ChooseListenerOrRuleAddTG
            action: aws:branch
            inputs:
              Choices:
                - NextStep: AddTGtoExistingRule
                  Variable: '{{ variable:IsDefaultRule }}'
                  StringEquals: 'False'
              Default: AddTGToExistingListener
          - name: AddTGToExistingListener
            action: aws:executeAwsApi
            nextStep: sleep
            isEnd: false
            inputs:
              ListenerArn: '{{ CalculateRuleSettings.ListenerArn }}'
              DefaultActions: '{{ CalculateRuleSettings.Actions }}'
              Service: elbv2
              Api: ModifyListener
          - name: AddTGtoExistingRule
            action: aws:executeAwsApi
            nextStep: sleep
            isEnd: false
            inputs:
              Actions: '{{ CalculateRuleSettings.Actions }}'
              RuleArn: '{{ RuleArn }}'
              Service: elbv2
              Api: ModifyRule
          - name: sleep
            action: aws:sleep
            nextStep: ChooseListenerOrRuleRollback
            isEnd: false
            onCancel: step:ChooseListenerOrRuleRollback
            onFailure: step:ChooseListenerOrRuleRollback
            inputs:
              Duration: PT{{ DurationSeconds }}S
          - name: ChooseListenerOrRuleRollback
            action: aws:branch
            inputs:
              Choices:
                - NextStep: RollbackTGSettings
                  Variable: '{{ variable:IsDefaultRule }}'
                  StringEquals: 'False'
              Default: RollbackListenerSettings
          - name: RollbackListenerSettings
            action: aws:executeAwsApi
            nextStep: RemovePermission
            isEnd: false
            inputs:
              ListenerArn: '{{ CalculateRuleSettings.ListenerArn }}'
              DefaultActions: '{{ CalculateRuleSettings.OriginalActions }}'
              Service: elbv2
              Api: ModifyListener
          - name: RollbackTGSettings
            action: aws:executeAwsApi
            nextStep: RemovePermission
            isCritical: false
            isEnd: false
            inputs:
              Actions: '{{ CalculateRuleSettings.OriginalActions }}'
              RuleArn: '{{ RuleArn }}'
              Service: elbv2
              Api: ModifyRule
          - name: RemovePermission
            action: aws:executeAwsApi
            nextStep: DeleteTargetGroup
            isEnd: false
            onFailure: step:DeleteTargetGroup
            inputs:
              FunctionName: '{{ LambdaFunctionName }}'
              StatementId: SID{{ automation:EXECUTION_ID }}
              Service: lambda
              Api: RemovePermission
          - description: Delete Target Group
            name: DeleteTargetGroup
            action: aws:executeAwsApi
            isEnd: true
            onFailure: Abort
            inputs:
              TargetGroupArn: '{{ CreateTargetGroup.TargetGroupArn }}'
              Service: elbv2
              Api: DeleteTargetGroup
  
  FISExperimentRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: fis.amazonaws.com
            Action: 'sts:AssumeRole'
            Condition:
              StringEquals:
                'aws:SourceAccount': !Sub '${AWS::AccountId}'
              ArnLike:
                'aws:SourceArn': !Sub 'arn:${AWS::Partition}:fis:${AWS::Region}:${AWS::AccountId}:experiment/*'
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSFaultInjectionSimulatorSSMAccess'
      Policies:
        - PolicyName: FISExperimentPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogDelivery
                  - logs:DescribeLogGroups
                  - logs:CreateLogStream
                  - logs:PutResourcePolicy
                  - logs:DescribeResourcePolicies
                  - logs:PutLogEvents
                Resource: '*'
              - Effect: Allow
                Action:
                  - elasticloadbalancing:DescribeRules
                  - elasticloadbalancing:ModifyRule
                  - elasticloadbalancing:ModifyListener
                  - elasticloadbalancing:CreateTargetGroup
                  - elasticloadbalancing:DeleteTargetGroup
                  - elasticloadbalancing:RegisterTargets
                Resource: '*'
              - Effect: Allow
                Action: elasticloadbalancing:AddTags
                Resource: !Sub 'arn:${AWS::Partition}:elasticloadbalancing:${AWS::Region}:${AWS::AccountId}:loadbalancer/${LoadBalancerName}'
              - Effect: Allow
                Action:
                  - lambda:AddPermission
                  - lambda:RemovePermission
                Resource: !Sub 'arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${AWS::StackName}-error-simulation'
      Path: '/service-role/'
   
  FISExperimentTemplate:
     Type: AWS::FIS::ExperimentTemplate
     Properties:
       Description: ALB-FIS-TEMPLATE
       Actions:
         SimulateFailureALB:
           ActionId: aws:ssm:start-automation-execution
           Parameters:
             documentArn: !Sub 'arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:document/${TargetGroupFailureSimulationDocument}'
             documentParameters: !Sub '{"LambdaFunctionName":["${ErrorSimulationFunction}"],"DurationSeconds":["${TestDurationInSeconds}"],"RuleArn":["${ListenerRuleArn}"],"FailurePercentage":["50"]}'
             maxDuration: PT62M # This value should be at least 2 minutes greater than the value DurationSeconds set for the SSM automation document to allow FIS to complete any rollback actions and finish the experiment.
       StopConditions:
         - Source: none
       Targets: {}
       RoleArn: !GetAtt FISExperimentRole.Arn
       Tags:
         Name: ALB-Failure-Simulation
       LogConfiguration:
         CloudWatchLogsConfiguration:
           LogGroupArn: !Sub 'arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:${FISExperimentLogGroup}:*'
         LogSchemaVersion: 2
       ExperimentOptions:
         AccountTargeting: single-account
         EmptyTargetResolutionMode: fail 

Outputs:
  ErrorSimulationFunctionArn:
    Description: ARN of the Error Simulation Lambda Function
    Value: !GetAtt ErrorSimulationFunction.Arn
  ErrorSimulationFunctionName:
    Description: Name of the Error Simulation Lambda Function
    Value: !Ref ErrorSimulationFunction
  ErrorSimulationFunctionRoleArn:
    Description: ARN of the IAM Role
    Value: !GetAtt ErrorSimulationFunctionRole.Arn
  SSMDocumentName:
    Description: Name of the SSM Automation Document
    Value: !Ref TargetGroupFailureSimulationDocument
