AWSTemplateFormatVersion: '2010-09-09'
Description: 'Add sidecar container to existing ECS task definition and update service'

Parameters:
  ExistingTaskDefinitionArn:
    Type: String
    Description: ARN of the existing ECS task definition

  SidecarContainerName:
    Type: String
    Description: Name of the sidecar container to add

  SidecarContainerImage:
    Type: String
    Description: Docker image for the sidecar container

  EcsClusterName:
    Type: String
    Description: Name of the ECS cluster

  EcsServiceName:
    Type: String
    Description: Name of the ECS service to update

  ControllerServerName:
    Type: String
    Description: Name of the Appsentinels controller

  ControllerPort:
    Type: String
    Description: Controller TCP port number

  EnvironmentType:
    Type: String
    Description: Type of environment (dev/test/prod)

  TapInterface:
    Type: String
    Description: Name of the interface to sniff on

  TapFilter:
    Type: String
    Description: BPF filter

  RelayProtocol:
    Type: String
    Description: Protocol to communicate with controller

  AwsRegion:
    Type: String
    Description: AWS region

Resources:
  UpdateTaskDefinitionFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.8
      Timeout: 300
      Code:
        ZipFile: |
          import boto3
          import cfnresponse

          def handler(event, context):
            try:
              ecs = boto3.client('ecs')
              
              existing_task_def_arn = event['ResourceProperties']['ExistingTaskDefinitionArn']
              sidecar_name = event['ResourceProperties']['SidecarContainerName']
              sidecar_image = event['ResourceProperties']['SidecarContainerImage']
              cluster_name = event['ResourceProperties']['EcsClusterName']
              service_name = event['ResourceProperties']['EcsServiceName']

              controller_name = event['ResourceProperties']['ControllerServerName']
              controller_port = event['ResourceProperties']['ControllerPort']
              environment_type = event['ResourceProperties']['EnvironmentType']
              tap_interface = event['ResourceProperties']['TapInterface']
              tap_filter = event['ResourceProperties']['TapFilter']
              relay_protocol = event['ResourceProperties']['RelayProtocol']
              aws_region = event['ResourceProperties']['AwsRegion']
              
              response = ecs.describe_task_definition(taskDefinition=existing_task_def_arn)
              task_def = response['taskDefinition']
              
              container_definitions = task_def['containerDefinitions']
              
              # Check if sidecar already exists
              sidecar_exists = any(container['name'] == sidecar_name for container in container_definitions)
              
              if not sidecar_exists:
                family = task_def['family']
                aws_group = '/ecs/' + family
                # Add new sidecar container
                new_container = {
                  'name': sidecar_name,
                  'image': sidecar_image,
                  'essential': True,
                  # Add other container properties as needed
                  'environment': [
                    {
                      'name': 'REMOTE_CONTROLLER_SERVER_NAME',
                      'value': controller_name
                    },
                    {
                      'name': 'REMOTE_CONTROLLER_SERVER_PORT',
                      'value': controller_port
                    },
                    {
                      'name': 'ENVIRONMENT',
                      'value': environment_type
                    },
                    {
                      'name': 'TAP_INTERFACE',
                      'value': tap_interface
                    },
                    {
                      'name': 'TAP_FILTER',
                      'value': tap_filter
                    },
                    {
                      'name': 'RELAY_PROTOCOL',
                      'value': relay_protocol
                    },
                    {
                      'name': 'TAP_PROFILE',
                      'value': 'low'
                    },
                    {
                      'name': 'MAX_SAAS_QUEUE_HOLDING_CAP',
                      'value': '4096'
                    },
                    {
                      'name': 'SAAS_CONNECTOR_LIMIT',
                      'value': '10'
                    }
                  ],
                  'logConfiguration': {
                    'logDriver': 'awslogs',
                    'options': {
                      'awslogs-group': aws_group,
                      'awslogs-create-group': 'true',
                      'awslogs-region': aws_region,
                      'awslogs-stream-prefix': 'ecs'
                    }
                  }
                }
                container_definitions.append(new_container)
              
              # Register new task definition
              new_task_def = ecs.register_task_definition(
                family=task_def['family'],
                taskRoleArn=task_def.get('taskRoleArn'),
                executionRoleArn=task_def.get('executionRoleArn'),
                networkMode=task_def['networkMode'],
                containerDefinitions=container_definitions,
                volumes=task_def.get('volumes', []),
                placementConstraints=task_def.get('placementConstraints', []),
                requiresCompatibilities=task_def.get('requiresCompatibilities', []),
                cpu=task_def.get('cpu'),
                memory=task_def.get('memory')
              )

              new_task_def_arn = new_task_def['taskDefinition']['taskDefinitionArn']
              
              # Update the ECS service
              ecs.update_service(
                cluster=cluster_name,
                service=service_name,
                taskDefinition=new_task_def_arn
              )
              
              cfnresponse.send(event, context, cfnresponse.SUCCESS, {'TaskDefinitionArn': new_task_def_arn})
              
            except Exception as e:
              print(str(e))
              cfnresponse.send(event, context, cfnresponse.FAILED, {'Error': str(e)})

  UpdateTaskDefinition:
    Type: Custom::UpdateTaskDefinition
    Properties:
      ServiceToken: !GetAtt UpdateTaskDefinitionFunction.Arn
      ExistingTaskDefinitionArn: !Ref ExistingTaskDefinitionArn
      SidecarContainerName: !Ref SidecarContainerName
      SidecarContainerImage: !Ref SidecarContainerImage
      EcsClusterName: !Ref EcsClusterName
      EcsServiceName: !Ref EcsServiceName
      ControllerServerName: !Ref ControllerServerName
      ControllerPort: !Ref ControllerPort
      EnvironmentType: !Ref EnvironmentType
      TapInterface: !Ref TapInterface
      TapFilter: !Ref TapFilter
      RelayProtocol: !Ref RelayProtocol
      AwsRegion: !Ref AwsRegion

  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      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: ECSAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - ecs:DescribeTaskDefinition
                  - ecs:RegisterTaskDefinition
                  - ecs:UpdateService
                Resource: '*'
              - Effect: Allow
                Action:
                  - iam:PassRole
                Resource:
                  - !Sub 'arn:aws:iam::${AWS::AccountId}:role/*TaskExecutionRole'
                  - !Sub 'arn:aws:iam::${AWS::AccountId}:role/*TaskRole'

Outputs:
  UpdatedTaskDefinitionArn:
    Description: ARN of the updated task definition
    Value: !GetAtt UpdateTaskDefinition.TaskDefinitionArn
