---
AWSTemplateFormatVersion: '2010-09-09'
Description: This template is to be deployed in the parent account, to allow cross-account creation of a patched test environment.
Parameters:
  ChildAccountId:
    Type: String
    Description: AWS Account ID of the child account (the account that AWS Systems
      Manager Automation will be targeting).
    MaxLength: 12
    MinLength: 12
    AllowedPattern: ^[0-9]{12}$
  KeyPairName:
    Type: String
    Description: Enter the name of the key pair that you use to connect to your EC2 instances
  SSMInstanceProfile:
    Type: String
    Description: Name of the Systems Manager instance profile attached to your EC2
      instances in the child account.
  KmsKeyArn:
    Type: String
    Description: (Optional) The Amazon Resource Name (ARN) of the CMK to use when encrypting lambda log group data.
Conditions:
  HasKeyArn: !Not [!Equals ["", !Ref 'KmsKeyArn']]
Resources:
  ParentAccountRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: AWS-SystemsManager-AutomationAdministrationRole-Parent
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: ssm.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: /
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonSSMAutomationRole
      Policies:
        - PolicyName: AssumeRole-AWSSystemsManagerAutomationExecutionRole
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - sts:AssumeRole
                Resource: !Sub 'arn:${AWS::Partition}:iam::*:role/AWS-SystemsManager-AutomationExecutionRole-Child'
              - Effect: Allow
                Action:
                  - organizations:ListAccountsForParent
                Resource:
                  - '*'
  AutomationDocument:
    Type: AWS::SSM::Document
    Properties:
      Content:
        description: This automation creates an AMI image, then launches an EC2 instance
          based on that AMI image.
        schemaVersion: '0.3'
        assumeRole: '{{AutomationAssumeRole}}'
        parameters:
          InstanceId:
            type: String
            description: (Required) The ID of the Amazon EC2 instance.
          NoReboot:
            type: Boolean
            default: false
            description: (Optional) Do not reboot the instance before creating the
              image.
          AutomationAssumeRole:
            type: String
            description: '(Optional) The ARN of the role that allows Automation to
              perform the actions on your behalf. '
            default: ''
          Workload:
            type: String
            description: '(Required) The workload this instance belongs to. '
          RoleARN:
            type: String
            description: (Required) The ARN of the lambda cross account access role.
        mainSteps:
          - name: createImage
            action: aws:createImage
            inputs:
              InstanceId: '{{InstanceId}}'
              ImageName: TEST_{{ Workload }}_ {{ InstanceId }}
              NoReboot: false
          - name: invokeLambda
            action: aws:invokeLambdaFunction
            inputs:
              FunctionName: !Sub 'arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:Automation-getInstanceTypeAndTags'
              Payload: '{"instanceId": "{{InstanceId}}", "ARN": "{{RoleARN}}"}'
          - name: invokeScript
            action: aws:executeScript
            outputs:
              - Name: instanceType
                Selector: $.Payload.instanceType
                Type: String
              - Name: instanceTags
                Selector: $.Payload.instanceTags
                Type: MapList
            inputs:
              Runtime: python3.6
              Handler: script_handler
              InputPayload:
                payload: '{{ invokeLambda.Payload }}'
                instanceId: '{{InstanceId}}'
              Script: |-
                import json
                def script_handler(event, context):
                  json_object = json.loads(event['payload'])
                  return {'instanceType': json_object['InstanceType'], 'instanceTags': json_object['Tags']}
          - name: launchInstance
            action: aws:runInstances
            inputs:
              ImageId: '{{ createImage.ImageId }}'
              IamInstanceProfileName: !Ref 'SSMInstanceProfile'
              InstanceType: '{{ invokeScript.instanceType }}'
              KeyName: !Ref 'KeyPairName'
              TagSpecifications:
                - ResourceType: instance
                  Tags:
                    - Key: Name
                      Value: TEST_{{ Workload }}_ {{ InstanceId }}
          - name: createTags
            action: aws:createTags
            maxAttempts: 3
            onFailure: Abort
            inputs:
              ResourceType: EC2
              ResourceIds: '{{ launchInstance.InstanceIds }}'
              Tags: '{{ invokeScript.instanceTags }}'
          - maxAttempts: 3
            inputs:
              Parameters:
                Operation: Install
              InstanceIds:
                - '{{ launchInstance.InstanceIds }}'
              DocumentName: AWS-RunPatchBaseline
            name: runPatchBaseline
            action: aws:runCommand
            timeoutSeconds: 800
            onFailure: Continue
      DocumentType: Automation
      Name: Automation-CreateTestEnvironment
  LambdaRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: getInstanceTypeAndTagsRole
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Path: /
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: AssumeRole-Lambda
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - sts:AssumeRole
                Resource: !Sub
                  - arn:${AWS::Partition}:iam::${childAccId}:role/lambda-cross-account-access-role
                  - childAccId: !Ref 'ChildAccountId'
  LambdaFunctionPermissions:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt 'LambdaFunctionGetInstanceTypeAndTags.Arn'
      Action: lambda:InvokeFunction
      Principal: !Ref 'ChildAccountId'
  LambdaFunctionLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: /aws/lambda/Automation-getInstanceTypeAndTags
      RetentionInDays: 7
      KmsKeyId: !If ['HasKeyArn', !Ref 'KmsKeyArn', !Ref 'AWS::NoValue']
  LambdaFunctionGetInstanceTypeAndTags:
    DependsOn:
      - LambdaFunctionLogGroup
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: Automation-getInstanceTypeAndTags
      Handler: index.handler
      Role: !GetAtt 'LambdaRole.Arn'
      Runtime: nodejs12.x
      Timeout: 300
      Code:
        ZipFile: |
          const aws = require("aws-sdk");
          exports.handler = async (event) => {
            console.log("Handling patched clones creation event: ", event);
            const sts = new aws.STS({ region: process.env.AWS_REGION });
            const stsParams = {
              RoleArn: event.ARN,
              RoleSessionName: "awsaccount_session",
            };

            try {
              const stsResults = await sts.assumeRole(stsParams).promise();
              const response = await buildResponse(stsResults, event.instanceId);
              console.log("Response: ", response);
              return response;
            } catch (err) {
              console.error(err);
              const message = `Error retrieving instance type and tags for instance id: ${event.instanceId}`;
              console.error(message);
              throw new Error(message);
            }
          };

          async function buildResponse(stsResults, instanceId) {
            const ec2Client = new aws.EC2({
              region: process.env.AWS_REGION,
              accessKeyId: stsResults.Credentials.AccessKeyId,
              secretAccessKey: stsResults.Credentials.SecretAccessKey,
              sessionToken: stsResults.Credentials.SessionToken,
            });

            const newTagObjects = await getEC2Tags(ec2Client, instanceId);
            const instanceType = await getEC2Type(ec2Client, instanceId);

            return {
              Tags: newTagObjects,
              InstanceType: instanceType.Value,
            };
          }

          async function getEC2Tags(ec2Client, instanceId) {
            const ec2Params = {
              Filters: [
                {
                  Name: "resource-id",
                  Values: [instanceId],
                },
              ],
            };

            const ec2TagsResponse = await ec2Client.describeTags(ec2Params).promise();

            const tagObjects = ec2TagsResponse.Tags;
            let newTagObjects = [];

            for (let tagObject of tagObjects) {
              const { ResourceId, ResourceType, ...newTagObject } = tagObject;
              if (newTagObject.Key === "Name") continue;
              newTagObjects.push(newTagObject);
            }
            return newTagObjects;
          }

          async function getEC2Type(ec2Client, instanceId) {
            const ec2InstanceType = await ec2Client
              .describeInstanceAttribute({
                Attribute: "instanceType",
                InstanceId: instanceId,
              })
              .promise();
            return ec2InstanceType.InstanceType;
          }


