# Unless explicitly stated otherwise, all files in this repository are
# licensed under the Apache 2 License.
#
# This product includes software developed at Datadog
# (https://www.datadoghq.com/). Copyright 2024 Datadog, Inc.
#
# This CloudFormation template creates a simple Auto Scaling group of Observability Pipelines
# Workers (OPWs), and is suitable for production workloads to trial the software.
AWSTemplateFormatVersion: 2010-09-09
Transform: AWS::LanguageExtensions
Description: Launches an Observability Pipelines Worker cluster to handle Datadog Agent workloads.
Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      -
        Label:
          default: "Observability Pipelines Worker Configuration"
        Parameters:
          - APIKey
          - PipelineID
          - DatadogSite
          - OPWEnv
          - OPWPort
      -
        Label:
          default: "Network Configuration"
        Parameters:
          - VPCID
          - SubnetIDs
      -
        Label:
          default: "EC2 Configuration"
        Parameters:
          - InstanceType
          - EBSSizeGB
      -
        Label:
          default: "Auto Scaling Group Configuration"
        Parameters:
          - ASGMinSize
          - ASGMaxSize
          - ASGCapacity
Parameters:
  APIKey:
    Description: Datadog API Key
    Type: String
    NoEcho: true
  PipelineID:
    Description: OPW Pipeline ID
    Type: String
  DatadogSite:
    Description: Datadog site to use. Defaults to US1.
    Type: String
    Default: "datadoghq.com"
    AllowedValues:
      - datadoghq.com
      - us3.datadoghq.com
      - us5.datadoghq.com
      - datadoghq.eu
      - ap1.datadoghq.com
  OPWEnv:
    Description: Environmental variables for Observability Pipelines Worker
    Type: String
    NoEcho: true
  OPWPort:
    Description: TCP port Observability Pipelines Worker will listen on
    Type: Number
    Default: 8282
  VPCID:
    Description: VPC to launch instances into.
    Type: "AWS::EC2::VPC::Id"
  SubnetIDs:
    Description: Subnets to launch instances into.
    Type: "List<AWS::EC2::Subnet::Id>"
  InstanceType:
    Description: EC2 instance type to launch. Allowed values include the AWS Graviton ARM-based types.
    Type: String
    Default: c6g.large
    AllowedValues:
      - t4g.nano
      - t4g.micro
      - t4g.small
      - t4g.medium
      - t4g.large
      - t4g.xlarge
      - t4g.2xlarge
      - c6g.medium
      - c6g.large
      - c6g.xlarge
      - c6g.2xlarge
      - c6g.4xlarge
      - c7g.medium
      - c7g.large
      - c7g.xlarge
      - c7g.2xlarge
      - c7g.4xlarge
  EBSSizeGB:
    Description: Storage to allocate to each instance for buffers, in GB.
    Type: Number
    Default: 288
  ASGMinSize:
    Description: Minimum instances to launch from the Auto Scaling group.
    Type: Number
    Default: 2
  ASGMaxSize:
    Description: Maximum instances to launch from the Auto Scaling group.
    Type: Number
    Default: 10
  ASGCapacity:
    Description: Desired instances to launch from the Auto Scaling group.
    Type: Number
    Default: 2

Mappings:
  # These are the supported regions. AMIs listed here are Ubuntu 22.04 LTS,
  # for ARM processors. If you need to add to this, you can find supported
  # images here: https://cloud-images.ubuntu.com/locator/ec2/
  AWSRegionToAMI:
    af-south-1:
      "HVMG2": "ami-0d39ef5932bf23109"
    ap-east-1:
      "HVMG2": "ami-021c24082c69e2867"
    ap-northeast-1:
      "HVMG2": "ami-02b8b87a5bfed4c7a"
    ap-northeast-2:
      "HVMG2": "ami-0ee413af3d419b913"
    ap-northeast-3:
      "HVMG2": "ami-0707b5724f3d370a1"
    ap-south-1:
      "HVMG2": "ami-0041facac80f93bbe"
    ap-south-2:
      "HVMG2": "ami-0544ba2d7ff852b4c"
    ap-southeast-1:
      "HVMG2": "ami-084594bed915f5e8f"
    ap-southeast-2:
      "HVMG2": "ami-07bf7b6acb5db5e13"
    ap-southeast-3:
      "HVMG2": "ami-07f185c39796fd830"
    ap-southeast-4:
      "HVMG2": "ami-0ce7ff6086582795a"
    ap-southeast-5:
      "HVMG2": "ami-0768220bae79b4901"
    ap-southeast-7:
      "HVMG2": "ami-0e4a91f0466e0c675"
    ca-central-1:
      "HVMG2": "ami-0d1bde699564f5a7a"
    ca-west-1:
      "HVMG2": "ami-0c62ba0d6cc805ca7"
    eu-central-1:
      "HVMG2": "ami-0bc586dd8476ba5b1"
    eu-central-2:
      "HVMG2": "ami-00f504206d607395a"
    eu-north-1:
      "HVMG2": "ami-008e58f8f6505bf76"
    eu-south-1:
      "HVMG2": "ami-0be337e76511c7b7b"
    eu-south-2:
      "HVMG2": "ami-0d777f40ee93f0622"
    eu-west-1:
      "HVMG2": "ami-04434961757c31b63"
    eu-west-2:
      "HVMG2": "ami-02c510094a29f0052"
    eu-west-3:
      "HVMG2": "ami-0bbeecfe5196eab29"
    il-central-1:
      "HVMG2": "ami-0a91f8186ca9ab853"
    me-central-1:
      "HVMG2": "ami-079335fbdd6702aa0"
    me-south-1:
      "HVMG2": "ami-0bb7bf7951ae5ff2c"
    mx-central-1:
      "HVMG2": "ami-09a423ffd5b0cbaaa"
    sa-east-1:
      "HVMG2": "ami-080924daa71cbdad8"
    us-east-1:
      "HVMG2": "ami-06daf9c2d2cf1cb37"
    us-east-2:
      "HVMG2": "ami-0edd45507c30e47d4"
    us-west-1:
      "HVMG2": "ami-01557579a54cccc40"
    us-west-2:
      "HVMG2": "ami-057a2512faa740640"

Resources:
  # Embedded security group allows access only to the load balancers, and only
  # on the Datadog Agent port.
  InstanceSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: "Allows access to OPW ports"
      VpcId: !Ref VPCID

  # Allow ingress only to the Datadog Agent port.
  InstanceSecurityGroupTCPIngress8282:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      IpProtocol: tcp
      FromPort: !Ref OPWPort
      ToPort: !Ref OPWPort
      CidrIp: 0.0.0.0/0
      GroupId:
        Fn::GetAtt:
          - InstanceSecurityGroup
          - GroupId

  # Allow egress to anywhere.
  InstanceSecurityGroupTCPEgress:
    Type: AWS::EC2::SecurityGroupEgress
    Properties:
      IpProtocol: tcp
      FromPort: 0
      ToPort: 65535
      CidrIp: 0.0.0.0/0
      GroupId:
        Fn::GetAtt:
          - InstanceSecurityGroup
          - GroupId

  # IAM Role allows SSH access over SSM, which can be useful for
  # running `observability-pipelines-worker tap`, or for other diagnostic purposes.
  IAMRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - ec2.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore"

  # This instance profile contains the SSM role above, and is useful
  # in case you need to add additional roles for other AWS services.
  IAMInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Roles:
        - !Ref IAMRole

  # Auto Scaling group handles scaling to meet application demand.
  LaunchTemplate:
    Type: AWS::EC2::LaunchTemplate
    Properties:
      LaunchTemplateName: !Sub ${AWS::StackName}-launch-template
      LaunchTemplateData:
        IamInstanceProfile:
          Name: !Ref IAMInstanceProfile
        InstanceType: !Ref InstanceType
        ImageId: !FindInMap [AWSRegionToAMI, !Ref "AWS::Region", HVMG2]
        TagSpecifications:
          - ResourceType: "instance"
            Tags:
              - Key: "Name"
                Value: !Sub ${AWS::StackName}-instance
        NetworkInterfaces:
          # OPW instances should not be directly accessible; the load balancer
          # should take care of that.
          - AssociatePublicIpAddress: false
            Groups:
              - !Ref InstanceSecurityGroup
            DeviceIndex: 0
        # NOTE: This does create an EBS drive for each instance launched as part
        # of the Auto Scaling group. However, it is deleted when the
        # instance is taken out of service. This makes this installation method
        # unsuitable for longer-term deployments, and may lead to data loss if
        # an instance is terminated.
        BlockDeviceMappings:
          - DeviceName: /dev/sdf
            Ebs:
              VolumeSize: !Ref EBSSizeGB
              VolumeType: gp3
              DeleteOnTermination: true
        # This user data script downloads and installs OPW, and sets it up
        # for Remote Configuration on the pipeline specified in the inputs.
        #
        # It also auto-mounts the EBS drive associated with this instance.
        UserData:
          Fn::Base64:
            !Sub |
            #!/bin/bash
            # Unless explicitly stated otherwise, all files in this repository are licensed under the Apache 2 License.
            #
            # This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2024 Datadog, Inc.
            get_opw_ebs_drive() {
                echo "/dev/nvme1n1"
            }
            
            sudo apt -y update
            sudo apt -y install apt-transport-https curl gnupg python3-pip
            
            # Install CloudFormation helper scripts
            sudo pip3 install https://s3.amazonaws.com/cloudformation-examples/aws-cfn-bootstrap-py3-latest.tar.gz
            sudo /usr/local/bin/cfn-init -v --stack ${AWS::StackName} --resource ASG --region ${AWS::Region}
            
            # Remove beta channel before GA
            sudo sh -c "echo 'deb [signed-by=/usr/share/keyrings/datadog-archive-keyring.gpg] https://apt.datadoghq.com/ stable observability-pipelines-worker-2' > /etc/apt/sources.list.d/datadog-observability-pipelines-worker.list"
            sudo touch /usr/share/keyrings/datadog-archive-keyring.gpg
            sudo chmod a+r /usr/share/keyrings/datadog-archive-keyring.gpg
            curl https://keys.datadoghq.com/DATADOG_APT_KEY_CURRENT.public | sudo gpg --no-default-keyring --keyring /usr/share/keyrings/datadog-archive-keyring.gpg --import --batch
            curl https://keys.datadoghq.com/DATADOG_APT_KEY_C0962C7D.public | sudo gpg --no-default-keyring --keyring /usr/share/keyrings/datadog-archive-keyring.gpg --import --batch
            curl https://keys.datadoghq.com/DATADOG_APT_KEY_F14F620E.public | sudo gpg --no-default-keyring --keyring /usr/share/keyrings/datadog-archive-keyring.gpg --import --batch
            
            sudo apt -y update
            sudo apt -y install observability-pipelines-worker datadog-signing-keys
            
            DD_OP_ENV_VARS_ENCODED="${OPWEnv}"
            DD_OP_ENV_VARS_DECODED="${!DD_OP_ENV_VARS_ENCODED//;/$'\n'}"
            
            sudo printf "DD_OP_API_ENABLED=true
            DD_SITE=${DatadogSite}
            DD_API_KEY=${APIKey}
            DD_OP_PIPELINE_ID=${PipelineID}
            $DD_OP_ENV_VARS_DECODED
            " > /etc/default/observability-pipelines-worker
            
            # Mount the associated EBS drive at the data directory.
            device=$(get_opw_ebs_drive)
            sudo mkfs.xfs $device || true
            sudo mount -o rw $device /var/lib/observability-pipelines-worker || true
            sudo chown observability-pipelines-worker:observability-pipelines-worker /var/lib/observability-pipelines-worker
            
            sudo systemctl restart observability-pipelines-worker
            sudo /usr/local/bin/cfn-signal -e $? --stack ${AWS::StackName} --resource ASG --region ${AWS::Region}
            
            # Uninstall CloudFormation helper scripts
            sudo pip3 uninstall -y aws-cfn-bootstrap-2.0
            sudo apt -y remove python3-pip
  ASG:
    Type: AWS::AutoScaling::AutoScalingGroup
    CreationPolicy:
      ResourceSignal:
        Timeout: PT15M
    Properties:
      MaxSize: !Ref ASGMaxSize
      MinSize: !Ref ASGMinSize
      LaunchTemplate:
        LaunchTemplateId: !Ref LaunchTemplate
        Version: !GetAtt LaunchTemplate.LatestVersionNumber
      DesiredCapacity: !Ref ASGCapacity
      VPCZoneIdentifier: !Ref SubnetIDs
      TargetGroupARNs:
        - !Ref NLBTarget8282

  # This load balancer should be used to communicate with the OPW
  # instances. It is set as internal-only, so it is not accessible
  # from the outside world.
  NLB:
    Type: AWS::ElasticLoadBalancingV2::LoadBalancer
    Properties:
      Type: network
      Scheme: internal
      Subnets: !Ref SubnetIDs
      SecurityGroups:
        - !Ref NLBSecurityGroup

  NLBListener8282:
    Type: AWS::ElasticLoadBalancingV2::Listener
    Properties:
      Protocol: TCP
      Port: !Ref OPWPort
      LoadBalancerArn: !Ref NLB
      DefaultActions:
        - Type: forward
          TargetGroupArn:
            Ref: NLBTarget8282

  NLBTarget8282:
    Type: AWS::ElasticLoadBalancingV2::TargetGroup
    Properties:
      Protocol: TCP
      Port: !Ref OPWPort
      VpcId: !Ref VPCID

  NLBSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: "Allows access to OPW ports from internal network"
      VpcId: !Ref VPCID

  # Though this is set as 0.0.0.0/0, since the load balancer is internal
  # it is not accessible from the outside world. This is just a
  # shorthand for "any part of your private network".
  NLBSecurityGroupTCPIngress82828:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      IpProtocol: tcp
      FromPort: !Ref OPWPort
      ToPort: !Ref OPWPort
      CidrIp: 0.0.0.0/0
      GroupId:
        Fn::GetAtt:
          - NLBSecurityGroup
          - GroupId

  # Allow egress from the load balancer to anywhere. Again, since the load
  # balancer is set internal-only, this is just a shorthand for "anywhere
  # in your private network".
  NLBSecurityGroupTCPEgress:
    Type: AWS::EC2::SecurityGroupEgress
    Properties:
      IpProtocol: tcp
      FromPort: 0
      ToPort: 65535
      CidrIp: 0.0.0.0/0
      GroupId:
        Fn::GetAtt:
          - NLBSecurityGroup
          - GroupId

Outputs:
  # Put this output into the `observability_pipelines_worker`
  # section of the Datadog Agent config.
  LoadBalancerDNS:
    Description: DNS to use instead of Splunk HEC address.
    Value: !GetAtt NLB.DNSName
  # If you need to attach additional roles to the instance, for example
  # to access S3, attach them to this profile.
  InstanceProfile:
    Description: Instance profile to attach additional roles to.
    Value: !Ref IAMInstanceProfile