
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0
---
AWSTemplateFormatVersion: "2010-09-09"
Description: Customizing Heuristics for OpsWorks
Parameters:
  ServerName:
    Type: String
    Description: OpsWorks CM Server Name
  OWCABucketName:
    Type: String
    Description: S3 Bucket Name for the OWCA server. Must have the format of aws-opsworks-cm-SERVER-NAME-string
    AllowedPattern: ^(aws-opsworks-cm-)[A-Za-z0-9\-]{1,46}[A-Za-z0-9]$
  RecipeBucketName:
    Type: String
    Description: S3 bucket name where the recipe archive is located. Needs to be in the same region as the Chef Automate Server.
  RecipeKeyName:
    Type: String
    Description: Full object key of the recipe to be executed at the end of the maintenance period. Needs to be a zip file and the recipe needs to be called "recipe.rb".
  NotificationEmail:
    Type: String
    Description: Email address to receive notifications on failed maintenance. Leave empty to disable.

Conditions:
  CreateSNSTopic: !Not [!Equals [ !Ref NotificationEmail, ""]]


Resources:
  SSMDoc:
    Type: "AWS::SSM::Document"
    Properties: 
      DocumentType: "Command"
      Content:
        schemaVersion: "2.2"
        description: Download and run 
        parameters:
          S3URL:
            type: String
            description: S3 Presigned URL
        mainSteps:
        - action: aws:runShellScript
          name: "CHOWDocument"
          inputs:
            runCommand:
            - curl -o /tmp/chow.zip "{{ S3URL }}"
            - sudo mkdir /tmp/chow
            - sudo unzip /tmp/chow.zip -d /tmp/chow/
            - sudo /opt/opscode/bin/chef-apply /tmp/chow/recipe.rb


  SNSTopic:
    Type: AWS::SNS::Topic
    Condition: CreateSNSTopic
    Properties:
      DisplayName: 'SNS Topic for CHOW'
      Subscription:
        -
          Endpoint:
            !Ref NotificationEmail
          Protocol: email

  CWESchedule:
    Type: AWS::Events::Rule
    Properties: 
      Description: Triggers the CHOW Lambda Function on a schedule
      ScheduleExpression: "rate(5 minutes)"
      State: "DISABLED"
      Targets:
        - 
          Arn: 
            !Sub arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:CHOW-${ServerName}
          Id: "TargetFunctionV1"

  PermissionForEventsToInvokeLambda: 
    Type: AWS::Lambda::Permission
    Properties: 
      FunctionName: 
        Ref: "LambdaFunction"
      Action: "lambda:InvokeFunction"
      Principal: "events.amazonaws.com"
      SourceArn: !GetAtt CWESchedule.Arn 

  PermissionForS3ToInvokeLambda: 
    Type: AWS::Lambda::Permission
    Properties: 
      FunctionName: 
        Ref: "LambdaFunction"
      Action: "lambda:InvokeFunction"
      Principal: "s3.amazonaws.com"
      SourceArn: !Sub arn:aws:s3:::${OWCABucketName}

  LambdaIAMRole:
    Type: "AWS::IAM::Role"
    Properties: 
      AssumeRolePolicyDocument: 
        Version: "2012-10-17"
        Statement: 
          - 
            Effect: "Allow"
            Principal: 
              Service: 
                - "lambda.amazonaws.com"
            Action: 
              - "sts:AssumeRole"
      Path: "/"

  LambdaIAMPolicy:
    Type: "AWS::IAM::Policy"
    Properties:
      PolicyName: "CHOWPolicy"
      Roles: 
        - 
          !Ref LambdaIAMRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
        - 
          Effect: Allow
          Action:
          - s3:GetObject
          - events:DescribeRule
          - events:EnableRule
          - events:DisableRule
          - cloudformation:DescribeStackResource
          - events:PutRule
          Resource:
          - !Sub arn:aws:cloudformation:${AWS::Region}:${AWS::AccountId}:stack/aws-opsworks-cm-instance-${ServerName}-*
          - !Sub arn:aws:s3:::${RecipeBucketName}/${RecipeKeyName}
          - !GetAtt CWESchedule.Arn
        - 
          Effect: Allow
          Action: ssm:SendCommand
          Resource: "arn:aws:ssm:*:*:document/*"
        - 
          Effect: Allow
          Action: ssm:SendCommand
          Resource: "arn:aws:ec2:*:*:instance/*"
        - 
          Effect: Allow
          Action: opsworks-cm:DescribeServers
          Resource: "*"
        - 
          Action:
            - 'logs:CreateLogGroup'
            - 'logs:CreateLogStream'
            - 'logs:PutLogEvents'
          Resource: 'arn:aws:logs:*:*:*'
          Effect: Allow

  LambdaSNSPermission:
    Condition: CreateSNSTopic
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: "CHOWSNSPolicy"
      Roles:
        - 
          !Ref LambdaIAMRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          Effect: Allow
          Action: SNS:Publish
          Resource:
            - !Ref SNSTopic

  S3CustomRole:
    Type: "AWS::IAM::Role"
    Properties: 
      AssumeRolePolicyDocument: 
        Version: "2012-10-17"
        Statement: 
          - 
            Effect: "Allow"
            Principal: 
              Service: 
                - "lambda.amazonaws.com"
            Action: 
              - "sts:AssumeRole"
      Path: "/"
      Policies:
        - PolicyName: LogsPermission
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - 
                Action:
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                Resource: 'arn:aws:logs:*:*:*'
                Effect: Allow
              - 
                Effect: Allow
                Action: s3:PutBucketNotification
                Resource: !Sub arn:aws:s3:::${OWCABucketName}


  S3CustomLambda:
    Type: "AWS::Lambda::Function"
    Properties:
      Description: "Creates a S3 object put trigger for Lambda"
      Handler: "index.lambda_handler"
      Role: !GetAtt S3CustomRole.Arn
      Runtime: "python2.7"
      Timeout: "5"
      Environment:
        Variables:
          "BUCKET": !Ref OWCABucketName
          "LAMBDA": !GetAtt LambdaFunction.Arn
          "OWCA": !Ref ServerName
      Code:
        ZipFile: |
          import boto3
          import os
          import cfnresponse

          BUCKET_NAME = os.environ['BUCKET']
          LAMBDA_FUNCTION = os.environ['LAMBDA']
          OWCA_SERVER_NAME = os.environ['OWCA']
          PREFIX = 'commands/' + OWCA_SERVER_NAME + '/start_maintenance'
          S3 = boto3.resource('s3')
          BUCKET_NOTIFICATION = S3.BucketNotification(BUCKET_NAME)

          DATA = {
              'LambdaFunctionConfigurations':
                  [
                      {
                          'LambdaFunctionArn': LAMBDA_FUNCTION,
                          'Events': ['s3:ObjectCreated:*'],
                          'Filter': {
                              'Key': {
                                  'FilterRules': [
                                      {
                                          'Name': 'prefix',
                                          'Value': PREFIX
                                      }
                                  ]
                              }
                          }
                      }
                  ]
          }


          def lambda_handler(event, context):
              if event['RequestType'] == 'Create':
                  print('Create Request by CFN...')
                  try:
                      BUCKET_NOTIFICATION.put(NotificationConfiguration=DATA)
                      rdata = {'Status': 'SUCCESS'}
                      cfnresponse.send(event, context, cfnresponse.SUCCESS, rdata)
                  except Exception as e:
                      print('Setting S3 trigger failed: {error!s}'.format(error=e))
                      rdata = {'Status': 'FAILURE'}
                      cfnresponse.send(event, context, cfnresponse.FAILURE, rdata)
              elif event['RequestType'] == 'Delete':
                  print('Delete Request by CFN...')
                  try:
                      BUCKET_NOTIFICATION.put(NotificationConfiguration={})
                      rdata = {'Status': 'SUCCESS'}
                      cfnresponse.send(event, context, cfnresponse.SUCCESS, rdata)
                  except Exception as e:
                      print('Removing S3 trigger failed: {error!s}'.format(error=e))
                      rdata = {'Status': 'FAILURE'}
                      cfnresponse.send(event, context, cfnresponse.FAILURE, rdata)
              elif event['RequestType'] == 'Update':
                  print('Update Request by CFN...')
                  try:
                      BUCKET_NOTIFICATION.put(NotificationConfiguration=DATA)
                      rdata = {'Status': 'SUCCESS'}
                      cfnresponse.send(event, context, cfnresponse.SUCCESS, rdata)
                  except Exception as e:
                      print('Setting S3 trigger failed: {error!s}'.format(error=e))
                      rdata = {'Status': 'FAILURE'}
                      cfnresponse.send(event, context, cfnresponse.FAILURE, rdata)




  S3LambdaTrigger:
    Type: "AWS::CloudFormation::CustomResource"
    Properties:
      ServiceToken: !GetAtt S3CustomLambda.Arn

  LambdaFunction:
    Type: "AWS::Lambda::Function"
    Properties:
      Description: "Executes SSM command on OWCA server after maintenance"
      FunctionName: !Sub CHOW-${ServerName}
      Handler: "index.lambda_handler"
      Role: !GetAtt LambdaIAMRole.Arn
      Runtime: "python2.7"
      Timeout: "5"
      Environment:
        Variables:
          "OWCASERVERNAME": !Ref ServerName
          "BUCKETNAME": !Ref RecipeBucketName
          "KEYNAME": !Ref RecipeKeyName
          "SSMDOC": !Ref SSMDoc
          "RULENAME": !Ref CWESchedule
          "SNSTOPIC": !If [ CreateSNSTopic, !Ref SNSTopic, ""]
      Code:
        ZipFile: |
          import boto3
          import os
          import sys
          from datetime import datetime
          from datetime import timedelta

          # Get the service clients.
          S3 = boto3.client('s3')
          SSM = boto3.client('ssm')
          CFN = boto3.client('cloudformation')
          OWCA = boto3.client('opsworkscm')
          EVENTS = boto3.client('events')
          SNS = boto3.client('sns')

          OWCA_SERVER_NAME = os.environ['OWCASERVERNAME']
          BUCKET_NAME = os.environ['BUCKETNAME']
          KEY_NAME = os.environ['KEYNAME']
          SSM_DOC = os.environ['SSMDOC']
          RULE_NAME = os.environ['RULENAME']
          SNS_TOPIC = os.environ['SNSTOPIC']


          def lambda_handler(event, context):
              # enable scheduler
              try:
                  if EVENTS.describe_rule(Name=RULE_NAME)['State'] == 'DISABLED':
                      set_rule_schedule(RULE_NAME)
              except Exception as e:
                  print('Enabling CWE rule failed: {error!s}'.format(error=e))
                  sys.exit(1)

              # Check if server is done with maintenance
              if get_owca_status() == 'UNDER_MAINTENANCE':
                  print('Server still under maintenance, waiting...')
                  return

              # validate if maintenance has finished
              elif get_owca_status() == 'HEALTHY':
                  # Generate the URL to get zip from bucket
                  try:
                      url = S3.generate_presigned_url(
                          ClientMethod='get_object',
                          Params={'Bucket': BUCKET_NAME, 'Key': KEY_NAME})
                  except Exception as e:
                      print('Generating S3 Presigned URL failed: {error!s}'.format(error=e))
                      sys.exit(1)
                  # get the new instance ID from the OWCA server
                  ec2id = get_instance(get_owca_stack())

                  # Send command
                  try:
                      SSM.send_command(
                          InstanceIds=[ec2id],
                          DocumentName=SSM_DOC,
                          Parameters={'S3URL': [url]})
                  except Exception as e:
                      print('Sending SSM Command failed: {error!s}'.format(error=e))
                      sys.exit(1)

                  # disable scheduler
                  try:
                      EVENTS.disable_rule(Name=RULE_NAME)
                  except Exception as e:
                      print('Disabling CWE rule failed: {error!s}'.format(error=e))
                      sys.exit(1)
              else:
                  error_message = (
                      'Server in {!s} state. Disabling CWE. Please verify cause via the OpsWorks for Chef Automate '
                      'console'.format(get_owca_status()))
                  print(error_message)
                  try:
                      EVENTS.disable_rule(Name=RULE_NAME)
                  except Exception as e:
                      print('Disabling CWE rule failed: {error!s}'.format(error=e))
                      sys.exit(1)
                  if SNS_TOPIC:
                      SNS.publish(TopicArn=SNS_TOPIC, Message=error_message)


          def get_instance(stack_arn):
              try:
                  instance_id = CFN.describe_stack_resource(
                      StackName=stack_arn,
                      LogicalResourceId='EC2Instance')
              except Exception as e:
                  print('Retrieving instance id failed: {error!s}'.format(error=e))
                  sys.exit(1)

              return instance_id['StackResourceDetail']['PhysicalResourceId']


          def get_owca_status():
              try:
                  instance_stack = OWCA.describe_servers(
                      ServerName=OWCA_SERVER_NAME
                  )
              except Exception as e:
                  print('Get OWCA server failed: {error!s}'.format(error=e))
                  sys.exit(1)

              return instance_stack['Servers'][0]['Status']


          def get_owca_stack():
              try:
                  instance_stack = OWCA.describe_servers(
                      ServerName=OWCA_SERVER_NAME)
              except Exception as e:
                  print('Getting OWCA CFN stack failed: {error!s}'.format(error=e))
                  sys.exit(1)

              return instance_stack['Servers'][0]['CloudFormationStackArn']


          def cron_gen():
              future_time = datetime.utcnow() + timedelta(minutes=20)
              return future_time.strftime('%M/1 %H * * ? *')


          def set_rule_schedule(rule):
              exp = 'cron(' + cron_gen() + ')'
              try:
                  EVENTS.put_rule(
                      Name=rule,
                      ScheduleExpression=exp,
                      State='ENABLED',
                      Description=EVENTS.describe_rule(Name=rule)['Description'])
              except Exception as e:
                  print('Setting CWE cron expression failed: {error!s}'.format(error=e))
                  EVENTS.disable_rule(Name=rule)
                  sys.exit(1)