AWSTemplateFormatVersion: '2010-09-09'
Parameters:
  ComplianceScanAutomationRoleName:
    Type: String
    Description: 'Name of the role that allows AWS Systems Manager Automation to perform required actions'
    Default: AutomationAssociationManagerRole
  CustomResourceLambdaFunctionName:
    Type: String
    Description: 'Lambda function name that creates the Systems Manager Association'
    Default: AutomationAssociationManager
  OSDetectionDocument:
    Type: String
    Description: 'Document name that performs AWS-RunInspecChecks execution by detecting OS type and version'
    Default: OS-Detection-Compliance-Scan
  SSMAssociationName:
    Type: String
    Description: 'Name of the Association that executes OS detection and resulting AWS-RunInspecChecks commands'
    Default: OS-Detection-Compliance-Scan-Association
  ProfileS3Bucket:
    Type: String
    Description: 'ARN of the bucket that houses Chef InSpec profiles'
    Default: 'arn:aws:s3:::aws-central-continuous-compliance-chef-inspec-code-east'
  InSpecProfile1S3Path:
    Type: String
    Description: "Specify the source path of the Inspec Profile. Format: {path: https://s3.amazonaws.com/${AWS::Partition}-s3BucketName/inspecProfileFolder}"
    Default: '{ "path": "https://s3.amazonaws.com/aws-central-continuous-compliance-chef-inspec-code/al2-static" }'
  InSpecProfile2S3Path:
    Type: String
    Description: "Specify the source path of the Inspec Profile. Format: {path: https://s3.amazonaws.com/${AWS::Partition}-s3BucketName/inspecProfileFolder}"
    Default: '{ "path": "https://s3.amazonaws.com/aws-central-continuous-compliance-chef-inspec-code/windows2016-static" }'
  ComplianceScanSchedule:
    Type: String
    Description: 'Schedule for how often compliance scans should run. Format cron(0 10 * * ? *) OR rate(7 days)'
    Default: cron(0 10 * * ? *)
Resources:
  SSMAutomationRole:
    Type: AWS::IAM::Role
    Properties: 
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
            - Effect: Allow
              Principal:
                Service: 
                  - ssm.amazonaws.com
                  - lambda.amazonaws.com
              Action: 'sts:AssumeRole'
      RoleName: !Ref ComplianceScanAutomationRoleName
  SSMAutomationPolicy: 
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Ref ComplianceScanAutomationRoleName
      Roles:
        - !Ref SSMAutomationRole
      PolicyDocument: 
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - 's3:GetObject'
              - 's3:ListBucket'
            Resource: 
              - !Sub '${ProfileS3Bucket}'
              - !Sub '${ProfileS3Bucket}/*'
          - Effect: Allow
            Action:
              - 'iam:CreateInstanceProfile'
              - 'iam:ListInstanceProfilesForRole'
              - 'iam:PassRole'
              - 'iam:GetInstanceProfile'
              - 'iam:AddRoleToInstanceProfile'
            Resource: '*'
          - Effect: Allow
            Action:
              - 'ec2:DescribeInstances'
              - 'ssm:DescribeInstanceInformation'
              - 'ssm:SendCommand'
            Resource: '*'
  CustomSSMDocument:
    Type: AWS::SSM::Document
    Properties:
      Content:
        description: >-
          Redirect executing Doc - AWS-RunInspecChecks on a group of managed instances using InSpec profile depending on the version of the OS.
        schemaVersion: '0.3'
        assumeRole: '{{AutomationAssumeRole}}'
        parameters:
          AutomationAssumeRole:
            default: ''
            type: String
            description: (Required) The ARN of the role that allows Automation to perform the actions on your behalf
          InstanceId:
            type: String
            description: (Required) EC2 Instance target
        mainSteps:
          - name: GetInstanceState
            action: 'aws:executeAwsApi'
            inputs:
              Service: ec2
              Api: DescribeInstances
              InstanceIds:
                - '{{InstanceId}}'
            outputs:
              - Name: runningstate
                Selector: '$.Reservations[0].Instances[0].State.Name'
                Type: String
          - name: DetermineIfRunning
            action: 'aws:branch'
            inputs:
              Choices:
                - NextStep: GetInstance
                  Variable: '{{GetInstanceState.runningstate}}'
                  StringEquals: running
              Default: Sleep
          - name: GetInstance
            action: 'aws:executeAwsApi'
            inputs:
              Service: ssm
              Api: DescribeInstanceInformation
              InstanceInformationFilterList:
                - key: InstanceIds
                  valueSet:
                    - '{{ InstanceId }}'
            outputs:
              - Name: platformType
                Selector: '$.InstanceInformationList[0].PlatformType'
                Type: String
              - Name: platformName
                Selector: '$.InstanceInformationList[0].PlatformName'
                Type: String
              - Name: platformVersion
                Selector: '$.InstanceInformationList[0].PlatformVersion'
                Type: String
          - name: ChooseOSforCommands
            action: 'aws:branch'
            inputs:
              Choices:
                - And:
                    - Variable: '{{GetInstance.platformName}}'
                      Contains: Amazon Linux
                    - Variable: '{{GetInstance.platformVersion}}'
                      Contains: '2'
                  NextStep: runInspecForOS1
                - And:
                    - Variable: '{{GetInstance.platformName}}'
                      Contains: Windows
                    - Variable: '{{GetInstance.platformVersion}}'
                      Contains: '2016'
                  NextStep: runInspecForOS2
              Default: Sleep
          - name: runInspecForOS1
            action: 'aws:executeAwsApi'
            maxAttempts: 3
            timeoutSeconds: 3600
            onFailure: Abort
            inputs:
              Service: ssm
              Api: SendCommand
              DocumentName: AWS-RunInspecChecks
              InstanceIds:
                - '{{ InstanceId }}'
              Parameters:
                sourceType:
                  - S3
                sourceInfo:
                  - !Ref InSpecProfile1S3Path
            isEnd: true
          - name: runInspecForOS2
            action: 'aws:executeAwsApi'
            maxAttempts: 3
            timeoutSeconds: 3600
            onFailure: Abort
            inputs:
              Service: ssm
              Api: SendCommand
              DocumentName: AWS-RunInspecChecks
              InstanceIds:
                - '{{ InstanceId }}'
              Parameters:
                sourceType:
                  - S3
                sourceInfo:
                  - !Ref InSpecProfile2S3Path
            isEnd: true
          - name: Sleep
            action: 'aws:sleep'
            inputs:
              Duration: PT3S
      DocumentType: Automation
      Name: !Ref OSDetectionDocument
  LambdaRoleForCustomCFResource:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
            - Effect: Allow
              Principal:
                Service: lambda.amazonaws.com
              Action: 'sts:AssumeRole'
      ManagedPolicyArns:
        - !Sub arn:${AWS::Partition}:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
  LambdaPolicyForCustomCFResource: 
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Roles:
        - !Ref LambdaRoleForCustomCFResource
      PolicyDocument: 
        Version: '2012-10-17'
        Statement:
          - Sid: SSMPermissions
            Effect: Allow
            Action: 
              - 'ssm:CreateAssociation'
              - 'ssm:ListAssociations'
              - 'ssm:DeleteAssociation'
              - 'ssm:UpdateAssociation'
              - 'iam:PassRole'
            Resource: '*'
  LambdaFunctionToCreateSSMAssociation:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Ref CustomResourceLambdaFunctionName
      Description: 'Creates an SSM Association as a Custom CloudFormation Resource to Execute Automation'
      Environment:
        Variables:
          ssmRole: !GetAtt SSMAutomationRole.Arn
          ssmDoc: !Ref OSDetectionDocument
          ssmAssociation: !Ref SSMAssociationName
          scanSchedule: !Ref ComplianceScanSchedule
      Handler: index.lambda_handler
      Role: !GetAtt LambdaRoleForCustomCFResource.Arn
      Runtime: python3.7 
      Timeout: 60
      Code:
        ZipFile: |
          import json
          import boto3
          import os 
          import cfnresponse
          
          ssmCustomDoc = os.environ['ssmDoc']
          ssmAutomationRole = os.environ['ssmRole']
          ssmAssociationName = os.environ['ssmAssociation']
          scanSchedule = os.environ['scanSchedule']

          ssm = boto3.client('ssm')

          def lambda_handler(event, context):
                 
            if event['RequestType'] == 'Create':
              response = ssm.create_association(
                Name=ssmCustomDoc,
                AssociationName=ssmAssociationName,
                Parameters={
                  'AutomationAssumeRole':[ssmAutomationRole]
                },
                Targets=[
                  {
                    'Key':'InstanceIds',
                    'Values':['*']
                  }
                ],
                AutomationTargetParameterName='InstanceId',
                ScheduleExpression=scanSchedule
              )

            if event['RequestType'] == 'Update' or event['RequestType'] == 'Delete':
              associations = ssm.list_associations(
              AssociationFilterList=[
                {
                  'key':'AssociationName',
                  'value':ssmAssociationName
                }
              ])
              associationID = associations['Associations'][0]['AssociationId']

              if event['RequestType'] == 'Update':
                response = ssm.update_association(
                  AssociationId=associationID,
                  Parameters={
                    'AutomationAssumeRole':[ssmAutomationRole]
                  },
                  Name=ssmCustomDoc,
                  Targets=[
                    {
                      'Key':'InstanceIds',
                      'Values':['*']
                    }
                  ],
                  AssociationName=ssmAssociationName,
                  AutomationTargetParameterName='InstanceId',
                  ScheduleExpression=scanSchedule
                )

              if event['RequestType'] == 'Delete':
                response = ssm.delete_association(
                  Name=ssmCustomDoc,
                  AssociationId=associationID
                )
            
            cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
  CustomSSMAssociation:
    Type: Custom::SSMAssociation
    Version: 1.0
    Properties:
      ServiceToken: !Sub ${LambdaFunctionToCreateSSMAssociation.Arn}
    DependsOn:
      - SSMAutomationPolicy
      - CustomSSMDocument
      - LambdaPolicyForCustomCFResource