AWSTemplateFormatVersion: 2010-09-09
Description: This template is used to build the vSensor AMI.(qs-1rmjte85h)

Metadata:
  QSLint:
    Exclusions: [ W9002, W9003, W9004, W9006 ]
  cfn-lint:
    config:
      ignore_checks:
        - E9101
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: Darktrace vSensor AMI Generation Parameters
        Parameters:
          - TargetAmiSSMName
          - ShortID
          - VsensorUpdatekey
          - LogGroupName
          - KeyPairName
          - VPCSubnetId
          - AMISecurityGroupId
    ParameterLabels:
      TargetAmiSSMName:
        default: Name for newly created AMI
      ShortID:
        default: Used to distinguish Amazon resources by name to allow multiple concurrent installations.
      VsensorUpdatekey:
        default: Update Key
      LogGroupName:
        default: CloudWatch Log Group
      LogGroupRetention:
        default: Log Group Retention Period
      KeyPairName:
        default: EC2 Key-pair Name
      VPCSubnetId:
        default: VPC Subnet ID for vSensor AMI generation
      AMISecurityGroupId:
        default: The security group to apply to the AMI EC2 Instance

Parameters:
  TargetAmiSSMName:
    Type: String
    Description: Name for newly created AMI
  ShortID:
    Type: String
    Description: Used to distinguish Amazon resources by name to allow multiple concurrent installations.
  VsensorUpdatekey:
    Type: String
    Description: Which Update Key should be used for the vSensor? (Contact your Darktrace Representative)
    Default: 'XXXXXX:XXXX'
    NoEcho: true
  LogGroupName:
    Type: String
    Description: The CloudWatch log group that the AMI creation process should log to.
    Default: Darktrace-vSensor-Quickstart
  LogGroupRetention:
    Type: Number
    AllowedValues:
      - 1
      - 3
      - 5
      - 7
      - 14
      - 30
      - 60
      - 90
      - 120
      - 150
      - 180
      - 365
      - 400
      - 545
      - 731
      - 1827
      - 3653
    Description: Number of days to retain Cloudwatch logs.
    Default: 30
  KeyPairName:
    Type: 'AWS::EC2::KeyPair::KeyName'
    Description: Which EC2 Key-pair should be used to connect to the vSensor?
  VPCSubnetId:
    Type: 'AWS::EC2::Subnet::Id'
    Description: Which VPC should the AMI EC2 Instance be created in?
  AMISecurityGroupId:
    Type: 'AWS::EC2::SecurityGroup::Id'
  FocalAmiId:
    Type: 'AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>'
    Default: '/aws/service/canonical/ubuntu/server-minimal/focal/stable/current/amd64/hvm/ebs-gp2/ami-id'


Resources:

# Create IAM Role for SSM Automation
  AutomationRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Action:
              - 'sts:AssumeRole'
            Principal:
              Service:
                - ssm.amazonaws.com
            Effect: Allow
        Version: 2012-10-17
      ManagedPolicyArns:
        - !Sub 'arn:${AWS::Partition}:iam::aws:policy/service-role/AmazonSSMAutomationRole'
      Policies:
        # ec2:Describe* cannot have resource specified.
        # See https://aws.amazon.com/blogs/security/demystifying-ec2-resource-level-permissions/
        - PolicyName: EC2Access
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: 'Allow'
                Action:
                  - ec2:DescribeInstances
                  - ec2:DescribeInstanceStatus
                Resource: '*'
        - PolicyName: SSMAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: 'Allow'
                Action:
                  - 'ssm:PutParameter'
                  - 'ssm:GetParameter'
                  - 'ssm:UpdateInstanceInformation'
                Resource:
                  - !Sub 'arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:parameter/*'
        - PolicyName: CWAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: 'Allow'
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                  - logs:DescribeLogGroups
                  - logs:DescribeLogStreams
                Resource:
                  - !Sub 'arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:*'
        - PolicyName: CFNAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: 'Allow'
                Action:
                  - cloudformation:SignalResource
                Resource:
                  - !Sub 'arn:${AWS::Partition}:cloudformation:${AWS::Region}:${AWS::AccountId}:stack/*/*'
      RoleName:
        !Sub
        - "${ShortID}-vSensor-AMI-SSM-Automation-IAM-Role"
        - ShortID: !Ref ShortID

# IAM Role AMI Instance
  AMIInstanceRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Action:
              - 'sts:AssumeRole'
            Principal:
              Service:
                - ec2.amazonaws.com
                - ssm.amazonaws.com
            Effect: Allow
        Version: 2012-10-17
      ManagedPolicyArns:
        - !Sub 'arn:${AWS::Partition}:iam::aws:policy/service-role/AmazonSSMAutomationRole'
      # https://docs.aws.amazon.com/service-authorization/latest/reference/list_awssystemsmanager.html
      Policies:
        - PolicyName: SSMAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: 'Allow'
                Action: ssm:StartAutomationExecution
                Resource:
                  - !Sub 'arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:document/*:VersionId}'
              - Effect: 'Allow'
                Action: ssm:GetAutomationExecution
                Resource:
                  - '*'
        - PolicyName: CWAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: 'Allow'
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                  - logs:DescribeLogGroups
                  - logs:DescribeLogStreams
                Resource:
                  - !Sub 'arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:*'
        - PolicyName: CleanAMIEC2Instance
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - 'ec2:StopInstances'
                  - 'ec2:TerminateInstances'
                Resource:
                  - !Sub 'arn:${AWS::Partition}:ec2:${AWS::Region}:${AWS::AccountId}:*'
        # https://docs.aws.amazon.com/service-authorization/latest/reference/list_amazonmessagedeliveryservice.html
        - PolicyName: EC2MessagesAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: 'Allow'
                Action:
                  - ec2messages:GetEndpoint
                  - ec2messages:GetMessages
                  - ec2messages:AcknowledgeMessage
                  - ec2messages:FailMessage
                  - ec2messages:SendReply
                Resource:
                  - '*'
        # https://docs.aws.amazon.com/service-authorization/latest/reference/list_amazonsessionmanagermessagegatewayservice.html
        - PolicyName: SSMMessagesAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: 'Allow'
                Action:
                  - ssmmessages:OpenControlChannel
                  - ssmmessages:CreateDataChannel
                  - ssmmessages:CreateControlChannel
                  - ssmmessages:OpenDataChannel
                Resource:
                  - '*'
      RoleName:
        !Sub
        - "${ShortID}-vSensor-AMI-IAM-Role"
        - ShortID: !Ref ShortID

  AMIInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Roles:
        - !Ref AMIInstanceRole
      InstanceProfileName:
        !Sub
        - "${ShortID}-vSensor-AMI-IAM-InstanceProfile"
        - ShortID: !Ref ShortID


  BuildEC2Instance:
    Type: AWS::EC2::Instance
    CreationPolicy: # This is to stop CFN from showing the instance as completed. The success signal is sent from SSM
      ResourceSignal:
        Timeout: PT30M
        Count: 1
    Properties:
      ImageId: !Ref FocalAmiId
      KeyName: !Ref KeyPairName
      IamInstanceProfile: !Ref AMIInstanceProfile
      InstanceType: "t3.medium"
      BlockDeviceMappings:
        - DeviceName: "/dev/sda1"
          Ebs:
            VolumeSize: 20
            Encrypted: true
            DeleteOnTermination: true
      Tags:
        - Key: Name
          Value: !Ref TargetAmiSSMName
      UserData: !Base64
          Fn::Join:
            - "\n"
            - - "#!/bin/bash"
              - "apt update"
              - "apt install -y awscli"
              - !Sub "aws ssm start-automation-execution --document-name ${SsmAMIBuild} --region ${AWS::Region}"
      SubnetId: !Ref VPCSubnetId
      SecurityGroupIds:
        - !Ref AMISecurityGroupId

  SsmAMIBuild:
    Metadata:
      cfn-lint:
        config:
          ignore_checks:
            - E9101
          ignore_reasons:
            E9101: Not applicable to actions like aws:executeAwsApi
    Type: AWS::SSM::Document
    Properties:
      Name:
        !Sub
          - "${ShortID}-vSensor-AMI-SSM"
          - ShortID: !Ref ShortID
      DocumentType: Automation
      Content:
        schemaVersion: "0.3"
        description: "Deploy and create AMI"
        assumeRole: !GetAtt AutomationRole.Arn
        parameters:
          TargetAmiSSMName:
            default: !Ref TargetAmiSSMName
            description: "SSM Parameter Name for the created AMI to pass to vSensor ASG."
            type: "String"
          StackName:
            default: !Sub "${AWS::StackName}"
            description: "Stack Name Input for cfn resource signal"
            type: "String"
          ShortID:
            default: !Ref ShortID
            description: "Used to set AMI name unique for this stack."
            type: "String"
          VsensorUpdatekey:
            default: !Ref VsensorUpdatekey
            description: Which Update Key should be used for the vSensor? (Contact your Darktrace Representative)
            type: String
          LogGroupName:
            default: !Ref LogGroupName
            description: The CloudWatch log group that the AMI creation process should log to.
            type: String

        mainSteps:
          - name: "InstanceIds"
            action: aws:executeAwsApi
            onFailure: "step:SignalFailure"
            nextStep: "InstallVsensor"
            inputs:
              Service: ec2
              Api: DescribeInstances
              Filters:
                - Name: "tag:Name"
                  Values: ["{{ShortID}}-vSensor-AMI"]
                - Name: "instance-state-name"
                  Values: ["running"]
            outputs:
              - Name: InstanceId
                Selector: "$.Reservations[0].Instances[0].InstanceId"
                Type: "String"

          - name: InstallVsensor
            action: aws:runCommand
            onFailure: "step:SignalFailure"
            nextStep: "ShutdownVSensor"
            timeoutSeconds: 3600
            inputs:
              DocumentName: AWS-RunShellScript
              InstanceIds:
                - "{{InstanceIds.InstanceId}}"
              Parameters:
                commands:
                  - '#!/bin/bash'
                  - 'bash <(wget -O - https://packages.darktrace.com/install) --updateKey {{VsensorUpdatekey}} || exit 1'
                  - 'wget https://s3.amazonaws.com/amazoncloudwatch-agent/ubuntu/amd64/latest/amazon-cloudwatch-agent.deb || exit 1'
                  - 'sudo dpkg -i -E ./amazon-cloudwatch-agent.deb || exit 1'
                  - 'systemctl enable amazon-cloudwatch-agent || exit 1'
                  - 'systemctl disable unattended-upgrades || exit 1' # Disable unattended upgrades so it doesn't interfere on boot. Re-enabled in userdata script.
                  - 'sleep 20' # Sleep just to make sure the machine is settled before letting SSM shut it down for snapshotting.
                  # If any steps fails signals CFN of Failure
              CloudWatchOutputConfig:
                CloudWatchOutputEnabled: "true"
                CloudWatchLogGroupName: "{{LogGroupName}}"

          # Shut the vSensor down before snapshot for AMI
          - name: "ShutdownVSensor"
            action: "aws:executeAwsApi"
            onFailure: "step:SignalFailure"
            nextStep: "CreateImage"
            inputs:
              Service: ec2
              Api: StopInstances
              InstanceIds:
                - "{{InstanceIds.InstanceId}}"

          # Create the custom AMI
          - name: CreateImage
            action: aws:createImage
            maxAttempts: 3
            onFailure: "step:SignalFailure"
            nextStep: "WritevSensorAmiId"
            inputs:
              InstanceId: "{{InstanceIds.InstanceId}}"
              ImageName: "{{ShortID}}-vSensor-AMI"
              NoReboot: true
            outputs:
              - Name: "vSensorAmiId"
                Selector: "ImageId"
                Type: "String"

          # Write the custom AMI ID to param store.
          - name: "WritevSensorAmiId"
            action: "aws:executeAwsApi"
            onFailure: "step:SignalFailure"
            nextStep: "TerminateAMIEC2Instance"
            inputs:
              Service: ssm
              Api: PutParameter
              Name: !Ref TargetAmiSSMName
              Type: "String"
              Overwrite: True
              Value: "{{CreateImage.vSensorAmiId}}"

          # Terminate the AMI EC2 Instance
          - name: "TerminateAMIEC2Instance"
            action: "aws:executeAwsApi"
            onFailure: "step:SignalFailure"
            nextStep: "SignalSuccess"
            inputs:
              Service: ec2
              Api: TerminateInstances
              InstanceIds: ["{{InstanceIds.InstanceId}}"]

          # If all steps complete successfully signals CFN of Success
          - name: "SignalSuccess"
            action: "aws:executeAwsApi"
            isEnd: True
            inputs:
              Service: cloudformation
              Api: SignalResource
              LogicalResourceId: "BuildEC2Instance"
              StackName: "{{StackName}}"
              Status: SUCCESS
              UniqueId: "{{InstanceIds.InstanceId}}"

          - name: "SignalFailure"
            action: "aws:executeAwsApi"
            inputs:
              Service: cloudformation
              Api: SignalResource
              LogicalResourceId: "BuildEC2Instance"
              StackName: "{{StackName}}"
              Status: FAILURE
              UniqueId: "{{InstanceIds.InstanceId}}"

  # From https://gist.github.com/drumadrian/e1601ab34e7f609b5075f65599108960
  # Deregister the AMI on Stack deletion to prevent failure.
  CleanupAMIOnDeleteLambdaRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - lambda.amazonaws.com
          Action:
          - sts:AssumeRole
      Policies:
        - PolicyName: AMICleanupLambdaPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
            - Effect: Allow
              Action:
                - ec2:DeregisterImage
              Resource:
                - !Sub 'arn:${AWS::Partition}:ec2:${AWS::Region}::image/*'
            - Effect: Allow
              Action:
                - ec2:DescribeImages
              Resource:
                - '*'
  CleanupAMIOnDeleteLambda:
    DependsOn: BuildEC2Instance
    Type: "AWS::Lambda::Function"
    Properties:
      Code:
        ZipFile:
         |
          #!/usr/bin/env python
          # -*- coding: utf-8 -*-
          import json
          import boto3
          import urllib3
          def deregister_ami(region, image_name):
              """
              Empties and deletes the AMI
              :param AMI_name:
              :return:
              """
              ec2_client = boto3.client('ec2', region_name=region)

              ami_id = None

              # Work out which AMI we're deleting.
              response = ec2_client.describe_images(
                  Filters=[
                      {
                          'Name': 'name',
                          'Values': [
                              image_name,
                          ]
                      },
                  ]
              )
              if len(response["Images"]) != 1:
                raise Exception("Couldn't find single image in: {}".format(str(images)))

              ami_id = response["Images"][0]["ImageId"]
              print("Found AMI ID to delete: {} ".format(ami_id))

              # Now deregister the AMI.
              try:
                response = ec2_client.deregister_image(
                    ImageId=ami_id
                )
              except Exception:
                print("Failed to deregister AMI ID: {}".format(ami_id))
                return

              print ("Successfully deleted AMI: {}".format(ami_id))

          def lambda_handler(event, context):
              try:
                  region = event['ResourceProperties']['Region']
                  parameter_name = event['ResourceProperties']['ImageName']
                  # Only deregister if we are deleting the custom resource (Cloudformation is deleting)
                  if event['RequestType'] == 'Delete':
                      deregister_ami(region, parameter_name)
                  else:
                      print("Skipping deregistering AMI, stack is not deleting.")

                  sendResponseCfn(event, context, "SUCCESS")
              except Exception as e:
                  print(e)
                  sendResponseCfn(event, context, "FAILED")

          def sendResponseCfn(event, context, responseStatus):
              response_body = {'Status': responseStatus,
                              'Reason': 'Log stream name: ' + context.log_stream_name,
                              'PhysicalResourceId': context.log_stream_name,
                              'StackId': event['StackId'],
                              'RequestId': event['RequestId'],
                              'LogicalResourceId': event['LogicalResourceId'],
                              'Data': json.loads("{}")}
              
              json_response_body = json.dumps(response_body)
              headers = {
                  'content-type' : '',
                  'content-length' : str(len(json_response_body))
              }

              try:
                  http = urllib3.PoolManager()
                  response = http.request('PUT', event['ResponseURL'], headers=headers, body=json_response_body)
                  print("Status code:", response.status)


              except Exception as e:
                  print("send(..) failed executing http.request(..):", e)

      Description: cleanup AMI on Delete Lambda function.
      FunctionName: !Sub
        - "${ShortID}-vSensor-AMI-Cleanup"
        - ShortID: !Ref ShortID
      Handler: index.lambda_handler
      Role : !GetAtt CleanupAMIOnDeleteLambdaRole.Arn
      Runtime: python3.8
      Timeout: 60
      Tags:
        - Key: darktrace-vsensor-quickstart
          Value: !Ref ShortID
      TracingConfig:
        Mode: Active

  CleanupAMIOnDeleteLambdaLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub "/aws/lambda/${CleanupAMIOnDeleteLambda}"
      RetentionInDays: !Ref LogGroupRetention
      
  CleanupAMIOnDeleteLogPermissions:
    Type: AWS::IAM::Policy
    Properties:
      Roles:
        - !Ref CleanupAMIOnDeleteLambdaRole
      PolicyName: !Sub "${ShortID}-vSensor-AMI-Cleanup-LambdaLogGroup"
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Action:
          - logs:CreateLogStream
          - logs:PutLogEvents
          Resource:
          - !Sub "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${CleanupAMIOnDeleteLambda}"
          - !Sub "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${CleanupAMIOnDeleteLambda}:*"
          - !Sub "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${CleanupAMIOnDeleteLambda}:*:*"

  CleanupAMIOnDelete:
    Type: AWS::CloudFormation::CustomResource
    Properties:
      ServiceToken: !GetAtt CleanupAMIOnDeleteLambda.Arn
      Region: !Sub "${AWS::Region}"
      ImageName: !Sub
        - "${ShortID}-vSensor-AMI"
        - ShortID: !Ref ShortID