AWSTemplateFormatVersion: '2010-09-09'
Description: >
  AWS CloudFormation template to create a new VPC
  or use an existing VPC for ECS deployment
  in Create Cluster Wizard
Parameters:
  EcsClusterName:
    Type: String
    Description: >
      Specifies the ECS Cluster Name with which the resources would be
      associated
    Default: 'demoCluster'
  KeyName:
    Type: String
    Description: >
      REQUIRED -  Specifies the name of an existing Amazon EC2 key pair
      to enable SSH access to the EC2 instances in your cluster.
  EcsAmiId:
    Type: String
    Description: REQUIRED - Default ECS Optimized AMI for us-west-2 region. Please change it to reflect your regions' latest ECS AMI-ID
    Default: 'ami-8e7bc4ee'
  EcsInstanceType:
    Type: String
    Description: >
      Specifies the EC2 instance type for your container instances.
      Defaults to t2.medium
    Default: t2.medium
    ConstraintDescription: must be a valid EC2 instance type.
  VpcId:
    Type: String
    Description: >
      Optional - Specifies the ID of an existing VPC in which to launch
      your container instances. If you specify a VPC ID, you must specify a list of
      existing subnets in that VPC. If you do not specify a VPC ID, a new VPC is created
      with atleast 1 subnet.
    Default: ''
    AllowedPattern: "^(?:vpc-[0-9a-f]{8}|)$"
    ConstraintDescription: >
      VPC Id must begin with 'vpc-' or leave blank to have a
      new VPC created
  SubnetIds:
    Type: CommaDelimitedList
    Description: >
      Optional - Specifies the Comma separated list of existing VPC Subnet
      Ids where ECS instances will run
    Default: ''
  SecurityGroupId:
    Type: String
    Description: >
      Optional - Specifies the Security Group Id of an existing Security
      Group. Leave blank to have a new Security Group created
    Default: ''
  VpcCidr:
    Type: String
    Description: Optional - Specifies the CIDR Block of VPC
    Default: '10.0.0.0/16'
  SubnetCidr1:
    Type: String
    Description: Specifies the CIDR Block of Subnet 1
    Default: '10.0.0.0/24'
  SubnetCidr2:
    Type: String
    Description: Specifies the CIDR Block of Subnet 2
    Default: '10.0.1.0/24'
  SubnetCidr3:
    Type: String
    Description: Specifies the CIDR Block of Subnet 3
    Default: ''
  IamRoleInstanceProfile:
    Type: String
    Description: >
      Specifies the Name or the Amazon Resource Name (ARN) of the instance
      profile associated with the IAM role for the instance
    Default: 'ecsInstanceRole'
  SecurityIngressFromPort:
    Type: Number
    Description: >
      Optional - Specifies the Start of Security Group port to open on
      ECS instances - defaults to port 0
    Default: '80'
  SecurityIngressToPort:
    Type: Number
    Description: >
      Optional - Specifies the End of Security Group port to open on ECS
      instances - defaults to port 65535
    Default: '80'
  SecurityIngressCidrIp:
    Type: String
    Description: >
      Optional - Specifies the CIDR/IP range for Security Ports - defaults
      to 0.0.0.0/0
    Default: 0.0.0.0/0
  VpcAvailabilityZones:
    Type: CommaDelimitedList
    Description: >
      Specifies a comma-separated list of 3 VPC Availability Zones for
      the creation of new subnets. These zones must have the available status.
    Default: 'us-west-2b,us-west-2c,us-west-2a'
  EbsVolumeSize:
    Type: Number
    Description: >
      Optional - Specifies the Size in GBs, of the newly created Amazon
      Elastic Block Store (Amazon EBS) volume
    Default: '22'
  EbsVolumeType:
    Type: String
    Description: Optional - Specifies the Type of (Amazon EBS) volume
    Default: 'gp2'
    AllowedValues:
      - ''
      - standard
      - io1
      - gp2
      - sc1
      - st1
    ConstraintDescription: Must be a valid EC2 volume type.
Conditions:
  CreateEC2LCWithKeyPair:
    !Not [!Equals [!Ref KeyName, '']]
  CreateNewSecurityGroup:
    !Equals [!Ref SecurityGroupId, '']
  CreateNewVpc:
    !Equals [!Ref VpcId, '']
  CreateSubnet1: !And
    - !Not [!Equals [!Ref SubnetCidr1, '']]
    - !Condition CreateNewVpc
  CreateSubnet2: !And
    - !Not [!Equals [!Ref SubnetCidr2, '']]
    - !Condition CreateSubnet1
  CreateSubnet3: !And
    - !Not [!Equals [!Ref SubnetCidr3, '']]
    - !Condition CreateSubnet2
  CreateEbsVolume: !And
    - !Not [!Equals [!Ref EbsVolumeSize, '0']]
    - !Not [!Equals [!Ref EbsVolumeType, '']]
Resources:
  MyCluster:
    Type: "AWS::ECS::Cluster"
    Properties:
      ClusterName: !Ref EcsClusterName
  Vpc:
    Condition: CreateSubnet1
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: !Ref VpcCidr
      EnableDnsSupport: 'true'
      EnableDnsHostnames: 'true'
  PubSubnetAz1:
    Condition: CreateSubnet1
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref Vpc
      CidrBlock: !Ref SubnetCidr1
      AvailabilityZone: !Select [ 0, !Ref VpcAvailabilityZones ]
  PubSubnetAz2:
    Condition: CreateSubnet2
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref Vpc
      CidrBlock: !Ref SubnetCidr2
      AvailabilityZone: !Select [ 1, !Ref VpcAvailabilityZones ]
  PubSubnetAz3:
    Condition: CreateSubnet3
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref Vpc
      CidrBlock: !Ref SubnetCidr3
      AvailabilityZone: !Select [ 2, !Ref VpcAvailabilityZones ]
  InternetGateway:
    Condition: CreateSubnet1
    Type: AWS::EC2::InternetGateway
  AttachGateway:
    Condition: CreateSubnet1
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId: !Ref Vpc
      InternetGatewayId: !Ref InternetGateway
  RouteViaIgw:
    Condition: CreateSubnet1
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref Vpc
  PublicRouteViaIgw:
    Condition: CreateSubnet1
    Type: AWS::EC2::Route
    DependsOn: AttachGateway
    Properties:
      RouteTableId: !Ref RouteViaIgw
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref InternetGateway
  PubSubnet1RouteTableAssociation:
    Condition: CreateSubnet1
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PubSubnetAz1
      RouteTableId: !Ref RouteViaIgw
  PubSubnet2RouteTableAssociation:
    Condition: CreateSubnet2
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PubSubnetAz2
      RouteTableId: !Ref RouteViaIgw
  PubSubnet3RouteTableAssociation:
    Condition: CreateSubnet3
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PubSubnetAz3
      RouteTableId: !Ref RouteViaIgw
  EcsSecurityGroup:
    Condition: CreateNewSecurityGroup
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: ECS Allowed Ports
      VpcId: !If [ CreateSubnet1, !Ref Vpc, !Ref VpcId ]
      SecurityGroupIngress:
        IpProtocol: tcp
        FromPort: !Ref SecurityIngressFromPort
        ToPort: !Ref SecurityIngressToPort
        CidrIp: !Ref SecurityIngressCidrIp
  EcsInstanceLc:
    Type: AWS::AutoScaling::LaunchConfiguration
    Properties:
      ImageId: !Ref EcsAmiId
      InstanceType: !Ref EcsInstanceType
      AssociatePublicIpAddress: true
      IamInstanceProfile: !Ref IamRoleInstanceProfile
      KeyName: !If [ CreateEC2LCWithKeyPair, !Ref KeyName, !Ref "AWS::NoValue" ]
      SecurityGroups: [ !If [ CreateNewSecurityGroup, !Ref EcsSecurityGroup, !Ref SecurityGroupId ] ]
      BlockDeviceMappings:
          - DeviceName: "/dev/sdm"
            Ebs:
              VolumeType: "gp2"
              Iops: "200"
              DeleteOnTermination: "true"
              VolumeSize: "20"
      UserData:
        Fn::Base64: !Sub |
           #!/bin/bash
           echo ECS_CLUSTER=${EcsClusterName} >> /etc/ecs/ecs.config
  EcsInstanceAsg:
    Type: AWS::AutoScaling::AutoScalingGroup
    Properties:
      VPCZoneIdentifier: !If
        - CreateSubnet1
        - !If
          - CreateSubnet2
          - !If
            - CreateSubnet3
            - [ !Sub "${PubSubnetAz1}, ${PubSubnetAz2}, ${PubSubnetAz3}" ]
            - [ !Sub "${PubSubnetAz1}, ${PubSubnetAz2}" ]
          - [ !Sub "${PubSubnetAz1}" ]
        - !Ref SubnetIds
      LaunchConfigurationName: !Ref EcsInstanceLc
      MinSize: '0'
      MaxSize:  '3'
      DesiredCapacity: '3'
      NotificationConfigurations:
      - TopicARN: !Ref ASGSNSTopic
        NotificationTypes:
        - autoscaling:EC2_INSTANCE_TERMINATE
      Tags:
        -
          Key: Name
          Value: !Sub "ECS Instance - ${AWS::StackName}"
          PropagateAtLaunch: 'true'
        -
          Key: Description
          Value: >
            This instance is the part of the Auto Scaling group which was created
            through ECS Console
          PropagateAtLaunch: 'true'
  taskdefinition:
    Type: "AWS::ECS::TaskDefinition"
    Properties:
      ContainerDefinitions:
        -
          Name: "ecs-sample-app"
          MountPoints:
            -
              SourceVolume: "my-vol"
              ContainerPath: "/var/www/my-vol"
          Image: "amazon/amazon-ecs-sample"
          Cpu: "10"
          PortMappings:
            -
              ContainerPort: "80"
              HostPort: "80"
          EntryPoint:
            - "/usr/sbin/apache2"
            - "-D"
            - "FOREGROUND"
          Memory: "500"
          Essential: "true"
        -
          Name: "busybox"
          Image: "busybox"
          Cpu: "10"
          EntryPoint:
            - "sh"
            - "-c"
          Memory: "500"
          Command:
            - "/bin/sh -c \"while true; do /bin/date > /var/www/my-vol/date; sleep 1; done\""
          Essential: "false"
          VolumesFrom:
            -
              SourceContainer: "ecs-sample-app"
      Volumes:
        -
          Host:
            SourcePath: "/var/lib/docker/vfs/dir/"
          Name: "my-vol"
  demoService:
    Type: "AWS::ECS::Service"
    Properties:
      Cluster: !Ref MyCluster
      DeploymentConfiguration:
        MaximumPercent: 100
        MinimumHealthyPercent: 50
      DesiredCount: 2
      TaskDefinition: !Ref taskdefinition
    DependsOn:
      - ECSServiceRole
      - taskdefinition
  ECSServiceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - ecs.amazonaws.com
          Action:
          - sts:AssumeRole
      Path: "/"
      Policies:
      - PolicyName: ecs-service
        PolicyDocument:
          Statement:
          - Effect: Allow
            Action:
            - elasticloadbalancing:Describe*
            - elasticloadbalancing:DeregisterInstancesFromLoadBalancer
            - elasticloadbalancing:RegisterInstancesWithLoadBalancer
            - ec2:Describe*
            - ec2:AuthorizeSecurityGroupIngress
            Resource: "*"
  SNSLambdaRole:
    Type: "AWS::IAM::Role"
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          -
            Effect: "Allow"
            Principal:
              Service:
                - "autoscaling.amazonaws.com"
            Action:
              - "sts:AssumeRole"
      ManagedPolicyArns:
      - arn:aws:iam::aws:policy/service-role/AutoScalingNotificationAccessRole
      Path: "/"
  LambdaExecutionRole:
    Type: "AWS::IAM::Role"
    Properties:
      Policies:
        -
          PolicyName: "lambda-inline"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                - autoscaling:CompleteLifecycleAction
                - logs:CreateLogGroup
                - logs:CreateLogStream
                - logs:PutLogEvents
                - ecs:ListContainerInstances
                - ecs:DescribeContainerInstances
                - ecs:UpdateContainerInstancesState
                - sns:Publish
                Resource: "*"
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          -
            Effect: "Allow"
            Principal:
              Service:
                - "lambda.amazonaws.com"
            Action:
              - "sts:AssumeRole"
      ManagedPolicyArns:
      - arn:aws:iam::aws:policy/service-role/AutoScalingNotificationAccessRole
      Path: "/"
  ASGSNSTopic:
    Type: "AWS::SNS::Topic"
    Properties:
      Subscription:
        -
          Endpoint:
             Fn::GetAtt:
                - "LambdaFunctionForASG"
                - "Arn"
          Protocol: "lambda"
    DependsOn: "LambdaFunctionForASG"
  LambdaFunctionForASG:
    Type: "AWS::Lambda::Function"
    Properties:
      Description: Gracefully drain ECS tasks from EC2 instances before the instances are
                   terminated by autoscaling.
      Handler: index.lambda_handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.6
      MemorySize: 128
      Timeout: 60
      Code:
        ZipFile: !Sub |
          import json
          import time
          import boto3

          CLUSTER = '${EcsClusterName}'
          REGION = '${AWS::Region}'

          ECS = boto3.client('ecs', region_name=REGION)
          ASG = boto3.client('autoscaling', region_name=REGION)
          SNS = boto3.client('sns', region_name=REGION)

          def find_ecs_instance_info(instance_id):
              paginator = ECS.get_paginator('list_container_instances')
              for list_resp in paginator.paginate(cluster=CLUSTER):
                  arns = list_resp['containerInstanceArns']
                  desc_resp = ECS.describe_container_instances(cluster=CLUSTER,
                                                               containerInstances=arns)
                  for container_instance in desc_resp['containerInstances']:
                      if container_instance['ec2InstanceId'] != instance_id:
                          continue

                      print('Found instance: id=%s, arn=%s, status=%s, runningTasksCount=%s' %
                            (instance_id, container_instance['containerInstanceArn'],
                             container_instance['status'], container_instance['runningTasksCount']))

                      return (container_instance['containerInstanceArn'],
                              container_instance['status'], container_instance['runningTasksCount'])

              return None, None, 0

          def instance_has_running_tasks(instance_id):
              (instance_arn, container_status, running_tasks) = find_ecs_instance_info(instance_id)
              if instance_arn is None:
                  print('Could not find instance ID %s. Letting autoscaling kill the instance.' %
                        (instance_id))
                  return False

              if container_status != 'DRAINING':
                  print('Setting container instance %s (%s) to DRAINING' %
                        (instance_id, instance_arn))
                  ECS.update_container_instances_state(cluster=CLUSTER,
                                                       containerInstances=[instance_arn],
                                                       status='DRAINING')

              return running_tasks > 0

          def lambda_handler(event, context):
              msg = json.loads(event['Records'][0]['Sns']['Message'])

              if 'LifecycleTransition' not in msg.keys() or \
                 msg['LifecycleTransition'].find('autoscaling:EC2_INSTANCE_TERMINATING') == -1:
                  print('Exiting since the lifecycle transition is not EC2_INSTANCE_TERMINATING.')
                  return

              if instance_has_running_tasks(msg['EC2InstanceId']):
                  print('Tasks are still running on instance %s; posting msg to SNS topic %s' %
                        (msg['EC2InstanceId'], event['Records'][0]['Sns']['TopicArn']))
                  time.sleep(5)
                  sns_resp = SNS.publish(TopicArn=event['Records'][0]['Sns']['TopicArn'],
                                         Message=json.dumps(msg),
                                         Subject='Publishing SNS msg to invoke Lambda again.')
                  print('Posted msg %s to SNS topic.' % (sns_resp['MessageId']))
              else:
                  print('No tasks are running on instance %s; setting lifecycle to complete' %
                        (msg['EC2InstanceId']))

                  ASG.complete_lifecycle_action(LifecycleHookName=msg['LifecycleHookName'],
                                                AutoScalingGroupName=msg['AutoScalingGroupName'],
                                                LifecycleActionResult='CONTINUE',
                                                InstanceId=msg['EC2InstanceId'])
  LambdaInvokePermission:
    Type: "AWS::Lambda::Permission"
    Properties:
       FunctionName: !Ref LambdaFunctionForASG
       Action: lambda:InvokeFunction
       Principal: "sns.amazonaws.com"
       SourceArn: !Ref ASGSNSTopic
  LambdaSubscriptionToSNSTopic:
    Type: AWS::SNS::Subscription
    Properties:
       Endpoint:
          Fn::GetAtt:
             - "LambdaFunctionForASG"
             - "Arn"
       Protocol: 'lambda'
       TopicArn: !Ref ASGSNSTopic
  ASGTerminateHook:
    Type: "AWS::AutoScaling::LifecycleHook"
    Properties:
      AutoScalingGroupName: !Ref EcsInstanceAsg
      DefaultResult: "ABANDON"
      HeartbeatTimeout: "900"
      LifecycleTransition: "autoscaling:EC2_INSTANCE_TERMINATING"
      NotificationTargetARN: !Ref ASGSNSTopic
      RoleARN:
         Fn::GetAtt:
         - "SNSLambdaRole"
         - "Arn"
    DependsOn: "ASGSNSTopic"
Outputs:
  EcsInstanceAsgName:
    Description: Auto Scaling Group Name for ECS Instances
    Value: !Ref EcsInstanceAsg
  SNSTopicForASG:
    Description: Topic used by ASG to send notifications when instance state is changing
    Value: !Ref ASGSNSTopic