# version: v1.0.0
AWSTemplateFormatVersion: '2010-09-09'
Description: Creates a Datadog Agentless Scanning delegate role
Parameters:
  AccountId:
    Type: String
    Description: The current AWS account ID. This parameter is for validation purposes only, and may be left empty.
    AllowedPattern: "|[0-9]{12}"
    Default: ""

  DatadogAPIKey:
    Type: String
    AllowedPattern: "[0-9a-f]{32}"
    Description: API key for the Datadog account
    NoEcho: true

  DatadogAPPKey:
    Type: String
    AllowedPattern: "[0-9a-f]{40}"
    Description: Application key for the Datadog account
    NoEcho: true

  DatadogSite:
    Type: String
    Description: The Datadog site to use for the Datadog Agentless Scanner
    Default: datadoghq.com
    AllowedValues:
      - datadoghq.com
      - datadoghq.eu
      - us3.datadoghq.com
      - us5.datadoghq.com
      - ap1.datadoghq.com
      - ap2.datadoghq.com

  AgentlessHostScanning:
    Type: String
    AllowedValues:
      - true
      - false
    Description: Enable Agentless Scanning of host vulnerabilities.
    Default: false

  AgentlessContainerScanning:
    Type: String
    AllowedValues:
      - true
      - false
    Description: Enable Agentless Scanning of container vulnerabilities.
    Default: false

  AgentlessLambdaScanning:
    Type: String
    AllowedValues:
      - true
      - false
    Description: Enable Agentless Scanning of Lambda vulnerabilities.
    Default: false

  AgentlessSensitiveDataScanning:
    Type: String
    AllowedValues:
      - true
      - false
    Description: Enable Agentless Scanning of datastores (S3 buckets).
    Default: false

  ScannerInstanceRoleARN:
    Type: CommaDelimitedList
    Description: The ARNs of the roles of the Datadog Agentless Scanner instances that will assume the delegate role.
    AllowedPattern: 'arn:aws:iam::[0-9]{12}:role/.*'

  ScannerDelegateRoleName:
    Type: String
    Description: The name of the role assumed by the Datadog Agentless Scanner
    Default: DatadogAgentlessScannerDelegateRole

Conditions:
  DSPMEnabled: !Equals
    - !Ref 'AgentlessSensitiveDataScanning'
    - 'true'

Rules:
  MustMatchAccountId:
    AssertDescription: 'Checking that AccountId matches the current AWS account ID'
    Assertions:
      - Assert: !Or [!Equals [!Ref AccountId, !Ref AWS::AccountId], !Equals [!Ref AccountId, ""]]
        AssertDescription: >-
          The current AWS account ID does not match the AWS account selected in Datadog.
          Please log in to the AWS account where you want to set up Datadog Agentless Scanning and try again.

Resources:
  ScannerDelegateRoleOrchestratorPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: Policy for the Datadog Agentless Scanner orchestrator allowing the creation and deletion of snapshots.
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: DatadogAgentlessScannerResourceTagging
            Action: 'ec2:CreateTags'
            Effect: Allow
            Resource:
              - 'arn:aws:ec2:*:*:volume/*'
              - 'arn:aws:ec2:*:*:snapshot/*'
              - 'arn:aws:ec2:*:*:image/*'
            Condition:
              StringEquals:
                'ec2:CreateAction':
                  - CreateSnapshot
                  - CreateVolume
                  - CopySnapshot
                  - CopyImage
          - Sid: DatadogAgentlessScannerVolumeSnapshotCreation
            Action: 'ec2:CreateSnapshot'
            Effect: Allow
            Resource: 'arn:aws:ec2:*:*:volume/*'
            Condition:
              StringNotEquals:
                'aws:ResourceTag/DatadogAgentlessScanner': 'false'
          - Sid: DatadogAgentlessScannerCopySnapshotSource
            Action: 'ec2:CopySnapshot'
            Effect: Allow
            Resource: 'arn:aws:ec2:*:*:snapshot/snap-*'
          - Sid: DatadogAgentlessScannerCopySnapshotDestination
            Action: 'ec2:CopySnapshot'
            Effect: Allow
            Resource: 'arn:aws:ec2:*:*:snapshot/${*}'
            Condition:
              'ForAllValues:StringLike':
                'aws:TagKeys': DatadogAgentlessScanner*
              StringEquals:
                'aws:RequestTag/DatadogAgentlessScanner': 'true'
          - Sid: DatadogAgentlessScannerSnapshotCreation
            Action: 'ec2:CreateSnapshot'
            Effect: Allow
            Resource: 'arn:aws:ec2:*:*:snapshot/*'
            Condition:
              'ForAllValues:StringLike':
                'aws:TagKeys': DatadogAgentlessScanner*
              StringEquals:
                'aws:RequestTag/DatadogAgentlessScanner': 'true'
          - Sid: DatadogAgentlessScannerSnapshotCleanup
            Action: 'ec2:DeleteSnapshot'
            Effect: Allow
            Resource: 'arn:aws:ec2:*:*:snapshot/*'
            Condition:
              StringEquals:
                'aws:ResourceTag/DatadogAgentlessScanner': 'true'
          - Sid: DatadogAgentlessScannerDescribeSnapshots
            Action: 'ec2:DescribeSnapshots'
            Effect: Allow
            Resource: '*'
          - Sid: DatadogAgentlessScannerEncryptedCopyGrant
            Action: 'kms:CreateGrant'
            Effect: Allow
            Resource: 'arn:aws:kms:*:*:key/*'
            Condition:
              'ForAnyValue:StringEquals':
                'kms:EncryptionContextKeys': 'aws:ebs:id'
              StringLike:
                'kms:ViaService': 'ec2.*.amazonaws.com'
              Bool:
                'kms:GrantIsForAWSResource': true
          - Sid: DatadogAgentlessScannerEncryptedCopyDescribe
            Action: 'kms:DescribeKey'
            Effect: Allow
            Resource: 'arn:aws:kms:*:*:key/*'
          - Sid: DatadogAgentlessScannerImageCleanup
            Action: 'ec2:DeregisterImage'
            Effect: Allow
            Resource: 'arn:aws:ec2:*:*:image/*'
            Condition:
              StringEquals:
                'aws:ResourceTag/DatadogAgentlessScanner': 'true'

  ScannerDelegateRoleWorkerPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: Policy for the Datadog Agentless Scanner worker allowing the listing and reading of snapshots.
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: DatadogAgentlessScannerSnapshotAccess
            Action:
              - 'ebs:ListSnapshotBlocks'
              - 'ebs:ListChangedBlocks'
              - 'ebs:GetSnapshotBlock'
            Effect: Allow
            Resource: 'arn:aws:ec2:*:*:snapshot/*'
            Condition:
              StringEquals:
                'aws:ResourceTag/DatadogAgentlessScanner': 'true'
          - Sid: DatadogAgentlessScannerDescribeSnapshots
            Action: 'ec2:DescribeSnapshots'
            Effect: Allow
            Resource: '*'
          - Sid: DatadogAgentlessScannerDescribeVolumes
            Action: 'ec2:DescribeVolumes'
            Effect: Allow
            Resource: '*'
          - Sid: DatadogAgentlessScannerDecryptEncryptedSnapshots
            Action: 'kms:Decrypt'
            Effect: Allow
            Resource: 'arn:aws:kms:*:*:key/*'
            Condition:
              'ForAnyValue:StringEquals':
                'kms:EncryptionContextKeys': 'aws:ebs:id'
              StringLike:
                'kms:ViaService': 'ec2.*.amazonaws.com'
          - Sid: DatadogAgentlessScannerKMSDescribe
            Action: 'kms:DescribeKey'
            Effect: Allow
            Resource: 'arn:aws:kms:*:*:key/*'
          - Sid: DatadogAgentlessScannerGetLambdaDetails
            Action: 'lambda:GetFunction'
            Effect: Allow
            Resource: 'arn:aws:lambda:*:*:function:*'
            Condition:
              StringNotEquals:
                'aws:ResourceTag/DatadogAgentlessScanner': 'false'
          - Sid: DatadogAgentlessScannerGetLambdaLayerDetails
            Action: 'lambda:GetLayerVersion'
            Effect: Allow
            Resource: 'arn:aws:lambda:*:*:layer:*:*'
            Condition:
              StringNotEquals:
                'aws:ResourceTag/DatadogAgentlessScanner': 'false'
          - Sid: DatadogAgentlessScannerECRAuthorizationToken
            Action:
              - "ecr:GetAuthorizationToken"
            Effect: Allow
            Resource: "*"
          - Sid: DatadogAgentlessScannerECRImages
            Action:
              - "ecr:GetDownloadUrlForLayer"
              - "ecr:BatchGetImage"
            Condition:
              StringNotEquals:
                "ecr:ResourceTag/DatadogAgentlessScanner": "false"
            Effect: Allow
            Resource: "arn:aws:ecr:*:*:repository/*"

  ScannerDelegateRoleWorkerDSPMPolicy:
    Type: AWS::IAM::ManagedPolicy
    Condition: DSPMEnabled
    Properties:
      Description: Policy for the Datadog Agentless Scanner worker allowing the listing and reading of S3 buckets.
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: DatadogAgentlessScannerAccessS3Objects
            Action: 's3:GetObject'
            Effect: Allow
            Resource: 'arn:aws:s3:::*/*'
          - Sid: DatadogAgentlessScannerListS3Buckets
            Action: 's3:ListBucket'
            Effect: Allow
            Resource: 'arn:aws:s3:::*'
          - Sid: DatadogAgentlessScannerDecryptS3Objects
            Action:
              - 'kms:Decrypt'
              - 'kms:GenerateDataKey'
            Effect: Allow
            Resource: 'arn:aws:kms:*:*:key/*'
            Condition:
              StringLike:
                'kms:ViaService': 's3.*.amazonaws.com'

  ScannerDelegateRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Ref 'ScannerDelegateRoleName'
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: EC2AssumeRole
            Effect: Allow
            Principal:
              AWS: '*'
            Condition:
              'ArnLike':
                'aws:PrincipalArn': !Ref 'ScannerInstanceRoleARN'
              StringEquals:
                'aws:PrincipalTag/Datadog': 'true'
                'aws:PrincipalTag/DatadogAgentlessScanner': 'true'
            Action: 'sts:AssumeRole'

      MaxSessionDuration: 3600
      ManagedPolicyArns:
        - !Ref 'ScannerDelegateRoleOrchestratorPolicy'
        - !Ref 'ScannerDelegateRoleWorkerPolicy'
        - !If [DSPMEnabled, !Ref 'ScannerDelegateRoleWorkerDSPMPolicy', !Ref 'AWS::NoValue']
      Description: Role assumed by the Datadog Agentless scanner agent to perform scans
      Tags:
        - Key: DatadogAgentlessScanner
          Value: 'true'
        - Key: Datadog
          Value: 'true'

  LambdaExecutionRoleDatadogAgentlessAPICall:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: "/"
      ManagedPolicyArns:
        - !Sub "arn:${AWS::Partition}:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"

  DatadogAgentlessAPICall:
    Type: "Custom::DatadogAgentlessAPICall"
    Properties:
      ServiceToken: !GetAtt "DatadogAgentlessAPICallFunction.Arn"
      TemplateVersion: "1.0.0"
      APIKey: !Ref "DatadogAPIKey"
      APPKey: !Ref "DatadogAPPKey"
      DatadogSite: !Ref "DatadogSite"
      AccountId: !Ref "AWS::AccountId"
      Hosts: !Ref "AgentlessHostScanning"
      Containers: !Ref "AgentlessContainerScanning"
      Lambdas: !Ref "AgentlessLambdaScanning"
      SensitiveData: !Ref "AgentlessSensitiveDataScanning"
      # Optional parameters
      DelegateRoleArn: !GetAtt "ScannerDelegateRole.Arn"
      OrchestratorPolicyArn: !Ref "ScannerDelegateRoleOrchestratorPolicy"
      WorkerPolicyArn: !Ref "ScannerDelegateRoleWorkerPolicy"
      WorkerDSPMPolicyArn: !If [DSPMEnabled, !Ref "ScannerDelegateRoleWorkerDSPMPolicy", !Ref "AWS::NoValue"]

  DatadogAgentlessAPICallFunction:
    Type: "AWS::Lambda::Function"
    Properties:
      Description: A function to call the Datadog Agentless API.
      Role: !GetAtt LambdaExecutionRoleDatadogAgentlessAPICall.Arn
      Handler: "index.handler"
      LoggingConfig:
        ApplicationLogLevel: "INFO"
        LogFormat: "JSON"
      Runtime: "python3.13"
      Timeout: 30
      Code:
        ZipFile: |
          #!/usr/bin/env python3
          
          import json
          import logging
          import signal
          from urllib.request import build_opener, HTTPHandler, HTTPError, Request
          import urllib.parse
          
          LOGGER = logging.getLogger()
          
          
          def call_datadog_agentless_api(event, method):
              template_version = event["ResourceProperties"]["TemplateVersion"]
              api_key = event["ResourceProperties"]["APIKey"]
              app_key = event["ResourceProperties"]["APPKey"]
              dd_site = event["ResourceProperties"]["DatadogSite"]
              account_id = event["ResourceProperties"]["AccountId"]
              hosts = event["ResourceProperties"]["Hosts"]
              containers = event["ResourceProperties"]["Containers"]
              lambdas = event["ResourceProperties"]["Lambdas"]
              sensitive_data = event["ResourceProperties"]["SensitiveData"]
              # Optional parameters
              launch_template_id = event["ResourceProperties"].get("LaunchTemplateId")
              asg_arn = event["ResourceProperties"].get("AutoScalingGroupArn")
              delegate_role_arn = event["ResourceProperties"].get("DelegateRoleArn")
              instance_role_arn = event["ResourceProperties"].get("InstanceRoleArn")
              instance_profile_arn = event["ResourceProperties"].get("InstanceProfileArn")
              orchestrator_policy_arn = event["ResourceProperties"].get("OrchestratorPolicyArn")
              worker_policy_arn = event["ResourceProperties"].get("WorkerPolicyArn")
              worker_dspm_policy_arn = event["ResourceProperties"].get("WorkerDSPMPolicyArn")
          
              # Make the url Request
              url = f"https://api.{dd_site}/api/v2/agentless_scanning/accounts/aws"
              headers = {
                  "DD-API-KEY": api_key,
                  "DD-APPLICATION-KEY": app_key,
                  "Dd-Call-Source": "cfn-agentless-quick-start",
                  "Dd-Installation-Version": template_version,
              }
          
              if method == "DELETE":
                  url = f"{url}/{account_id}"
                  request = Request(url, headers=headers, method="DELETE")
                  try:
                      return urllib.request.urlopen(request)
                  except HTTPError as e:
                      if e.status < 500:
                          # For most client errors, the best option is to continue with the
                          # stack deletion, since users have no way to fix the request, and
                          # at least this way they can clean up the scanner resources.
                          return e
                      else:
                          raise
          
              elif method == "POST":
                  values = {
                      "meta": {
                          "installation_mode": "cloudformation",
                          "installation_version": template_version,
                          "cloudformation_stack_id": event["StackId"],
                          "resources": {
                              "launch_template_id": launch_template_id,
                              "asg_arn": asg_arn,
                              "delegate_role_arn": delegate_role_arn,
                              "instance_role_arn": instance_role_arn,
                              "instance_profile_arn": instance_profile_arn,
                              "orchestrator_policy_arn": orchestrator_policy_arn,
                              "worker_policy_arn": worker_policy_arn,
                              "worker_dspm_policy_arn": worker_dspm_policy_arn,
                          },
                      },
                      "data": {
                          "id": account_id,
                          "type": "aws_scan_options",
                          "attributes": {
                              "vuln_containers_os": containers == "true",
                              "vuln_host_os": hosts == "true",
                              "lambda": lambdas == "true",
                              "sensitive_data": sensitive_data == "true",
                          },
                      },
                  }
                  data = json.dumps(values)
                  data = data.encode("utf-8")  # data should be bytes
                  url_account = f"{url}/{account_id}"
                  if is_agentless_scanning_enabled(url_account, headers):
                      request = Request(url_account, data=data, headers=headers, method="PATCH")
                  else:
                      request = Request(url, data=data, headers=headers, method="POST")
                  request.add_header("Content-Type", "application/vnd.api+json; charset=utf-8")
                  request.add_header("Content-Length", len(data))
                  response = urllib.request.urlopen(request)
                  return response
              else:
                  LOGGER.error("Unsupported HTTP method.")
                  return None
          
          
          def is_agentless_scanning_enabled(url_account, headers):
              """Check if agentless scanning is already enabled for the account"""
              request = Request(url_account, headers=headers, method="GET")
              try:
                  urllib.request.urlopen(request)
              except HTTPError as e:
                  if e.status == 404:
                      return False
                  else:
                      raise
              return True
          
          
          def handler(event, context):
              """Handle Lambda event from AWS"""
              try:
                  if event["RequestType"] == "Create":
                      LOGGER.info("Received Create request.")
                      response = call_datadog_agentless_api(event, "POST")
                      send_response(
                          event,
                          context,
                          "SUCCESS",
                          {
                              "Message": f"Datadog Agentless Scanning activated (status: {response.status}).",
                          },
                      )
                  elif event["RequestType"] == "Update":
                      LOGGER.info("Received Update request.")
                      send_response(
                          event,
                          context,
                          "SUCCESS",
                          {"Message": "Update not supported, no operation performed."},
                      )
                  elif event["RequestType"] == "Delete":
                      LOGGER.info("Received Delete request.")
                      response = call_datadog_agentless_api(event, "DELETE")
                      send_response(
                          event,
                          context,
                          "SUCCESS",
                          {
                              "Message": f"Datadog Agentless Scanning deactivated (status: {response.status}).",
                          },
                      )
                  else:
                      LOGGER.error(
                          "Failed - received unexpected request: %s", event["RequestType"]
                      )
                      send_response(
                          event,
                          context,
                          "FAILED",
                          {"Message": "Unexpected event received from CloudFormation"},
                      )
              except Exception as e:  # pylint: disable=W0702
                  LOGGER.exception("Failed - exception thrown during processing.")
                  send_response(
                      event,
                      context,
                      "FAILED",
                      {"Message": f"Exception during processing: {e}"},
                  )
          
          
          def send_response(event, context, response_status, response_data):
              """Send a resource manipulation status response to CloudFormation"""
              response_body = json.dumps(
                  {
                      "Status": response_status,
                      "Reason": "See the details in CloudWatch Log Stream: "
                      + context.log_stream_name,
                      "PhysicalResourceId": event.get(
                          "PhysicalResourceId", context.invoked_function_arn
                      ),
                      "StackId": event["StackId"],
                      "RequestId": event["RequestId"],
                      "LogicalResourceId": event["LogicalResourceId"],
                      "Data": response_data,
                  }
              )
              formatted_response = response_body.encode("utf-8")
          
              LOGGER.info("ResponseURL: %s", event["ResponseURL"])
              LOGGER.info("ResponseBody: %s", response_body)
          
              opener = build_opener(HTTPHandler)
              request = Request(event["ResponseURL"], data=formatted_response, method="PUT")
              request.add_header("Content-Type", "application/json; charset=utf-8")
              request.add_header("Content-Length", len(formatted_response))
              response = opener.open(request)
              LOGGER.info("Status code: %s", response.status)
              LOGGER.info("Status message: %s", response.msg)
          
          
          def timeout_handler(_signal, _frame):
              """Handle SIGALRM"""
              raise Exception("Time exceeded")
          
          
          signal.signal(signal.SIGALRM, timeout_handler)

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: "Required"
        Parameters:
          - ScannerInstanceRoleARN
          - ScannerDelegateRoleName
      - Label:
          default: "Advanced"
        Parameters:
          - AgentlessSensitiveDataScanning
          - AccountId

