SSM Document for attaching extension to a Lambda function and enable fault injection

attach_chaos_extension_ssm_document.yaml

schemaVersion: "0.3"
description: "Enable fault injection (Chaos) to a given lambda function by attaching a Lambda extension"
assumeRole: "{{ AutomationAssumeRole }}"
parameters:
  FunctionName:
    type: "String"
    description: "The name of the lambda function to add the layer"
  LayerArn:
    type: "String"
    description: "The ARN of the layer to add to the lambda function"
  AliasName:
    type: "String"
    description: "(Optional) The name of the lambda alias needs to be updated."
    default: ""
  AutomationAssumeRole:
    type: "String"
    description: "The ARN of the role that allows Automation to perform the actions on your behalf."
  ChaosMode:
    type: "String"
    description: "(Optional) Whether we want to enable or disable chaos. (ENABLED (default) or DISABLED)."
    allowedValues: [
      "ENABLED",
      "DISABLED"
    ]
    default: "ENABLED"
mainSteps:
  - name: addLayer
    description: Change a given Lambda configuration by adding a layer
    action: "aws:executeScript"
    inputs:
      Runtime: python3.8
      Handler: handler
      InputPayload:
        FunctionName: "{{FunctionName}}"
        LayerArn: "{{LayerArn}}"
        AliasName: "{{AliasName}}"
        ChaosMode: "{{ChaosMode}}"
      Script: |-
        from typing import Dict
        import boto3
        
        
        CHAOS_ENV_VARS = {
            'AWS_LAMBDA_EXEC_WRAPPER': '/opt/bootstrap',
            'CHAOS_EXTENSION__LAMBDA__ENABLE_LATENCY': 'true',
            'CHAOS_EXTENSION__LAMBDA__LATENCY_VALUE': '60',
            'CHAOS_EXTENSION__LAMBDA__LATENCY_PROBABILITY': '1',
        }
        
        
        class DocumentConfig:
            def __init__(self, function_name, layer_arn, alias_name, chaos_mode):
                self.function_name = function_name
                self.layer_arn = layer_arn
                self.alias_name = alias_name
                self.chaos_mode = chaos_mode
        
            @classmethod
            def from_event(cls, event):
                return cls(event['FunctionName'], event['LayerArn'], event['AliasName'], event['ChaosMode'])
        
        
        client = boto3.client('lambda')
        

        def handler(event, context):
            # Parse the input event to an internal class
            config: DocumentConfig = DocumentConfig.from_event(event)

            if config.chaos_mode == 'ENABLED':
              enable_chaos(config)
            else:
              disable_chaos(config)


        def enable_chaos(config):            
        
            lambda_config_existing = client.get_function_configuration(FunctionName=config.function_name)
        
            updated_layers = _get_updated_layers(lambda_config_existing, config.layer_arn)
            updated_variables = _add_chaos_vars(lambda_config_existing)
        
            _update_function(config, updated_layers, updated_variables)


        def disable_chaos(config):            
        
            lambda_config_existing = client.get_function_configuration(FunctionName=config.function_name)
        
            updated_layers = _get_updated_layers(lambda_config_existing, config.layer_arn, include_chaos_layer=False)
            updated_variables = _remove_chaos_vars(lambda_config_existing)
        
            _update_function(config, updated_layers, updated_variables)                    
        
        
        def _update_function(config, updated_layers, updated_variables):
            update_response = client.update_function_configuration(
                FunctionName=config.function_name,
                Environment={
                    'Variables': updated_variables,
                },
                Layers=updated_layers
            )
            # If response is not successful, raise an exception
            if update_response['ResponseMetadata']['HTTPStatusCode'] != 200:
                raise Exception(f'Failed to update function configuration for {config.function_name}')

            if config.alias_name:
              # Update the alias to point to the new updated version
              alias_response = client.update_alias(FunctionName=config.function_name, Name=config.alias_name, FunctionVersion=update_response['Version'])

              # If response is not successful, raise an exception
              if alias_response['ResponseMetadata']['HTTPStatusCode'] != 200:
                  raise Exception(f'Failed to update function alias {config.alias_name}')
        
        
        def _get_updated_layers(lambda_config_existing, layer_arn, include_chaos_layer=True):
            existing_layers = lambda_config_existing['Layers']
        
            # Get existing layer ARNs (without the layer we want to add if already connected)
            new_layers = [layer['Arn'] for layer in existing_layers if layer['Arn'] != layer_arn]
        
            # Add the new layer to the list of layers
            if include_chaos_layer:
                new_layers.append(layer_arn)
        
            return new_layers
        
        
        def _add_chaos_vars(lambda_config_existing: Dict[str, Dict[str, str]]) -> Dict[str, str]:
            existing_variables = lambda_config_existing['Environment']['Variables']
        
            updated_variables = {**existing_variables, **CHAOS_ENV_VARS}
            return updated_variables
        
        
        def _remove_chaos_vars(lambda_config_existing: Dict[str, Dict[str, str]]) -> Dict[str, str]:
            existing_variables: Dict[str, str] = lambda_config_existing['Environment']['Variables']
        
            updated_variables = {item: value for (item, value) in existing_variables.items() if item not in CHAOS_ENV_VARS}
            return updated_variables