#------------------------------------------------------------------------------
#
# Template: AWSCloudFormationInstanceAutomation.yaml
# Purpose:  Manage Automation for instances
#
# TO DO:
#   [ ] Lambda Function
#   [ ] Permission for automate documents on managed instances.
#   [ ] SNS notification for success & failure.
#------------------------------------------------------------------------------
AWSTemplateFormatVersion: 2010-09-09
Description: Manage Automation for instances

Parameters:
  Schedule:
    Type: String
    Description: Schedule for Owner for Maintenance Window (Either cron or Rate expression)
    Default: cron(0 4 ? * SUN *)

  TimeZone:
    Type: String
    Description: Time zone for scheduled maintenance window executions in Internet Assigned Numbers Authority (IANA) format.
    Default: US/Eastern

  Duration:
    Type: String
    Description: Duration for Maintenance Window automation job in hours
    Default: 5

  OwnerInformation:
    Type: String
    Description: User-provided value that will be included in any CloudWatch events
    Default: Patch Management Team

#-----------------------------------------------------------
# Resources
#-----------------------------------------------------------
Resources:
 SSMDocument:
   Type: AWS::SSM::Document
   Properties:
     DocumentType:  Automation
     Name: CustomAutomationDocument
     Content:
       schemaVersion: '0.3'
       assumeRole: '{{ AutomationAssumeRole }}'
       description: 'Run the RunPatchBaseline document on instances with RebootOption RebootIfNeeded option'
       parameters:
         InstanceIds:
           type: StringList
           description: (Required) EC2 Instance(s) to restart
         AutomationAssumeRole:
           type: String
           description: (Optional) The ARN of the role that allows Automation to perform the actions on your behalf.
           default: ''
         WorkingDirectory:
           default: ''
           description: (Optional) The path to the working directory on your instance.
           type: String
           maxChars: 4096
         ExecutionTimeout:
           default: '3600'
           description: (Optional) The time in seconds for a command to complete before it is considered to have failed. Default is 3600 (1 hour). Maximum is 172800 (48 hours).
           type: String
         Operation:
           type: String
           description: (Required) The update or configuration to perform on the instance. The system checks if patches specified in the patch baseline are installed on the instance. The install operation installs patches missing from the baseline.
           allowedValues:
             - Scan
             - Install
         SnapshotId:
           type: String
           description: (Optional) The snapshot ID to use to retrieve a patch baseline snapshot.
           allowedPattern: '(^$)|^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$'
           default: ''
         InstallOverrideList:
           type: String
           description: (Optional) An https URL or an Amazon S3 path-style URL to the list of patches to be installed. This patch installation list overrides the patches specified by the default patch baseline.
           allowedPattern: '(^$)|^https://.+$|^s3://([^/]+)/(.*?([^/]+))$'
           default: ''
         RebootOption:
           type: String
           description: '(Optional) Reboot behavior after a patch Install operation. If you choose NoReboot and patches are installed, the instance is marked as non-compliant until a subsequent reboot and scan.'
           allowedValues:
             - RebootIfNeeded
             - NoReboot
           default: RebootIfNeeded
       mainSteps:
         - name: RunPatchBaseline
           timeoutSeconds: 120
           action: 'aws:runCommand'
           maxAttempts: 3
           inputs:
             Parameters:
               SnapshotId: '{{ SnapshotId }}'
               InstallOverrideList: '{{ InstallOverrideList }}'
               Operation: '{{ Operation }}'
             InstanceIds: '{{ InstanceIds }}'
             DocumentName: AWS-RunPatchBaseline

 MaintenanceWindow:
   Type: AWS::SSM::MaintenanceWindow
   Properties:
     Description: Maintenance Window to trigger the automation
     AllowUnassociatedTargets: false
     Cutoff: 1
     Duration: !Ref Duration
     Name: UpdateSSMAgentMaintenanceWindow
     Schedule: !Ref Schedule
     ScheduleTimezone: !Ref TimeZone

 MaintenanceWindowTarget:
   Type: AWS::SSM::MaintenanceWindowTarget
   Properties:
     Description: Target for Maintenance Window
     Name: MaintenanceWindowTarget
     OwnerInformation: !Ref OwnerInformation
     ResourceType: INSTANCE
     Targets:
       - Key: tag:target
         Values:
           - DummyTarget
     WindowId: !Ref MaintenanceWindow

 MaintenanceWindowTask:
   Type: AWS::SSM::MaintenanceWindowTask
   Properties:
     Description: Apply automation to instance across regions.
     Name: LambdaTask
     MaxConcurrency: 1
     MaxErrors: 1
     Priority: 0
     TaskType: LAMBDA
     TaskArn: !Ref LambdaForAutomationFunction
     WindowId: !Ref MaintenanceWindow
     Targets:
       - Key: WindowTargetIds
         Values:
           - !Ref MaintenanceWindowTarget

 AWSSystemsManagerAutomationExecutionRole:
   Type: 'AWS::IAM::Role'
   Properties:
     RoleName: AWS-SystemsManager-AutomationExecutionRole
     AssumeRolePolicyDocument:
       Version: 2012-10-17
       Statement:
         - Effect: Allow
           Principal:
             AWS: !Ref AWS::AccountId
           Action:
             - 'sts:AssumeRole'
         - Effect: Allow
           Principal:
             Service: ssm.amazonaws.com
           Action:
             - 'sts:AssumeRole'
     ManagedPolicyArns:
       - 'arn:aws:iam::aws:policy/service-role/AmazonSSMAutomationRole'
     Path: /
     Policies:
       - PolicyName: ExecutionPolicy
         PolicyDocument:
           Version: 2012-10-17
           Statement:
             - Effect: Allow
               Action:
                 - 'resource-groups:ListGroupResources'
                 - 'tag:GetResources'
               Resource: '*'
             - Effect: Allow
               Action:
                 - 'iam:PassRole'
               Resource: !Sub >-
                 arn:${AWS::Partition}:iam::${AWS::AccountId}:role/AWS-SystemsManager-AutomationExecutionRole

 AWSSystemsManagerAutomationAdministrationRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: AWS-SystemsManager-AutomationAdministrationRole
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: ssm.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      Policies:
        - PolicyName: AssumeRole-AWSSystemsManagerAutomationExecutionRole
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - sts:AssumeRole
                Resource: !Sub arn:${AWS::Partition}:iam::${AWS::AccountId}:role/AWS-SystemsManager-AutomationExecutionRole
              - Effect: Allow
                Action:
                  - 'organizations:ListAccountsForParent'
                Resource: '*'

 LambdaForAutomationFunction:
   Type: AWS::Lambda::Function
   Properties:
     FunctionName: SSM-AutomationHandler
     Description:
       Automation the instances on Maintenance Window Event
     Handler: index.lambda_handler
     Role: !GetAtt LambdaAutomationIAMRole.Arn
     Runtime: python3.7
     MemorySize: 128
     Timeout: 600
     Environment:
       Variables:
         ASSUME_EXECUTION_ROLE: 'AWS-SystemsManager-AutomationExecutionRole'
     Code:
       ZipFile:
         Fn::Sub: |
            import boto3
            import os
            import json

            # SSM target execution Role
            ASSUME_EXECUTION_ROLE = os.environ['ASSUME_EXECUTION_ROLE']

            def lambda_handler(event, context):
                print('Event : ', event)

                document_version = ''
                document_parameters = ''

                if [x for x in (event.get('Document'), event.get('Targets')) if x is None]:
                    print('No Document/Targets information present in the event: ', event)
                    return {
                        'statusCode': 400,
                        'body': json.dumps('No Document/Targets information present in the event')
                }
                document = event['Document']
                targets = event['Targets']
                window_id = event['WindowId']
                task_execution_id = event['TaskExecutionId']
                document_name = document['Name']
                document_version = document['Version']

                if not document.get('Parameters'):
                    print('Document parameters are missing in event: ', event)
                    return {
                        'statusCode': 400,
                        'body': json.dumps('No Document/Targets information present in the event')
                }
                document_parameters = document['Parameters']
                target_parameter_name = event['TargetParameterName']

                account_id = context.invoked_function_arn.split(":")[4]
                target_locations = list()
                target_location = {}
                regions = list()
                target_location['Accounts'] = [account_id]
                target_location['TargetLocationMaxConcurrency'] = event['MaxConcurrency']
                target_location['TargetLocationMaxErrors'] = event['MaxErrors']
                target_location['ExecutionRoleName'] = ASSUME_EXECUTION_ROLE
                for region_id in event['Regions']:
                    resource_group_client = boto3.client('resource-groups', region_name=region_id)
                    try:
                        instances = resource_group_client.list_group_resources(
                                        GroupName=targets[0]['Values'][0]
                                    )
                        if instances.get('ResourceIdentifiers') and len(instances['ResourceIdentifiers']) > 0:
                            regions.append(region_id)
                    except Exception as e:
                        print('No instances under account {}, region {}'.format(account_id, region_id))
                        continue
                if len(regions) > 0:
                    target_location['Regions'] = regions
                    target_locations.append(target_location)
                    print('Running patching on target_locations {} '.format(target_locations))
                    client = boto3.client('ssm')
                    try:
                        response = client.start_automation_execution (
                            DocumentName=document_name,
                            DocumentVersion=document_version,
                            Parameters=document_parameters,
                            TargetParameterName=target_parameter_name,
                            Targets = targets,
                            MaxConcurrency= event['MaxConcurrency'],
                            MaxErrors=event['MaxErrors'],
                            TargetLocations=target_locations
                        )
                        print('AutomationExecutionId triggered for the job: ', response['AutomationExecutionId'])
                    except Exception as e:
                        print('Automation failed for WindowId: {}, TaskExecutionId: {} with exception {}'.format(window_id, task_execution_id, str(e)))
                        return {
                            'statusCode': 400,
                            'body': json.dumps('Automation failed for WindowId: {}, TaskExecutionId: {} with exception {}'.format(window_id, task_execution_id, str(e)))
                        }

 LambdaAutomationIAMRole:
   Type: AWS::IAM::Role
   Properties:
     AssumeRolePolicyDocument:
       Version: 2012-10-17
       Statement:
         - Effect: Allow
           Principal:
             Service:
               - lambda.amazonaws.com
           Action:
             - sts:AssumeRole
     Path: /
     Policies:
      - PolicyName: LambdaExecutionPolicy
        PolicyDocument:
          Version: 2012-10-17
          Statement:
            - Effect: Allow
              Action:
                - logs:CreateLogGroup
                - logs:CreateLogStream
                - logs:PutLogEvents
              Resource: arn:aws:logs:*:*:log-group:*
            - Effect: Allow
              Action:
                - iam:GetRole
                - iam:PassRole
              Resource: !GetAtt AWSSystemsManagerAutomationAdministrationRole.Arn
            - Effect: Allow
              Action:
                - resource-groups:List*
                - resource-groups:Get*
                - tag:GetResources
                - tag:GetTagValues
                - tag:DescribeReportCreation
                - tag:GetTagKeys
              Resource: "*"
            - Effect: Allow
              Action:
                - ssm:StartAutomationExecution
              Resource: !Sub arn:aws:ssm:${AWS::Region}:${AWS::AccountId}:automation-definition/CustomAutomationDocument:1
