---
AWSTemplateFormatVersion: '2010-09-09'
Description: Creates a Cloud9 IDE for the EKS workshop
Parameters:
  EksWorkshopC9InstanceType:
    Description: Cloud9 instance type
    Type: String
    Default: t3.small
    AllowedValues:
      - t2.micro
      - t3.micro
      - t3.small
      - t3.medium
    ConstraintDescription: Must be a valid Cloud9 instance type
  EksWorkshopC9EnvType: 
    Description: Environment type.
    Default: self
    Type: String
    AllowedValues: 
      - self
      - 3rdParty
    ConstraintDescription: must specify self or 3rdParty.
  WorkshopOwnerArn: 
    Type: String
    Description: The Arn of the Cloud9 Owner to be set if 3rdParty deployment.
    Default: ""
  EksWorkshopC9InstanceVolumeSize: 
    Type: Number
    Description: The Size in GB of the Cloud9 Instance Volume. 
    Default: 30
  RepositoryOwner: 
    Type: String
    Description: The owner of the GitHub repository to be used to bootstrap Cloud9
    Default: "aws-samples"
  RepositoryName: 
    Type: String
    Description: The name of the GitHub repository to be used to bootstrap Cloud9
    Default: "eks-workshop-v2"
  RepositoryRef: 
    Type: String
    Description: The Git reference to be used to bootstrap Cloud9
    Default: "main"
  Cloud9Name: 
    Type: String
    Description: Name of the Cloud9 instance
    Default: "none"
  Cloud9Subnet:
    Type: String
    Description: If you want Cloud9 in a specific subnet
    Default: "none"
  ResourcesPrecreated: 
    Type: String
    Description: Whether lab infrastructure has been pre-provisioned
    Default: "false"
    AllowedValues: 
      - "false"
      - "true"

Conditions: 
  Create3rdPartyResources: !Equals [ !Ref EksWorkshopC9EnvType, 3rdParty ]
  IsCloud9NotNamed: !Equals [ !Ref Cloud9Name, none ]
  IsCloud9SubnetSpecified: !Not [ !Equals [ !Ref Cloud9Subnet, none ]]

Resources:
  EksWorkshopC9Role:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - ec2.amazonaws.com
            - ssm.amazonaws.com
          Action:
          - sts:AssumeRole
      ManagedPolicyArns:
      - arn:aws:iam::aws:policy/AdministratorAccess
      - arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
      Path: "/"

  EksWorkshopC9LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - lambda.amazonaws.com
          Action:
          - sts:AssumeRole
      Path: "/"
      Policies:
      - PolicyName:
          Fn::Join:
          - ''
          - - EksWorkshopC9LambdaPolicy-
            - Ref: AWS::Region
        PolicyDocument:
          Version: '2012-10-17'
          Statement:
          - Effect: Allow
            Action:
            - logs:CreateLogGroup
            - logs:CreateLogStream
            - logs:PutLogEvents
            Resource: arn:aws:logs:*:*:*
          - Effect: Allow
            Action:
            - cloudformation:DescribeStacks
            - cloudformation:DescribeStackEvents
            - cloudformation:DescribeStackResource
            - cloudformation:DescribeStackResources
            - ec2:DescribeInstances
            - ec2:AssociateIamInstanceProfile
            - ec2:ModifyInstanceAttribute
            - ec2:ReplaceIamInstanceProfileAssociation
            - ec2:DescribeIamInstanceProfileAssociations
            - ec2:DescribeVolumes
            - ec2:ModifyVolume
            - ec2:DescribeVolumesModifications
            - ec2:RebootInstances
            - iam:ListInstanceProfiles
            - iam:PassRole
            - ssm:DescribeInstanceInformation
            - ssm:SendCommand
            - ssm:GetCommandInvocation
            Resource: "*"

  EksWorkshopC9BootstrapInstanceLambda:
    Type: Custom::EksWorkshopC9BootstrapInstanceLambda
    DependsOn:
    - EksWorkshopC9LambdaExecutionRole
    Properties:
      ServiceToken:
        Fn::GetAtt:
        - EksWorkshopC9BootstrapInstanceLambdaFunction
        - Arn
      REGION:
        Ref: AWS::Region
      Cloud9Name: !GetAtt EksWorkshopC9Instance.Name
      EnvironmentId:
        Ref: EksWorkshopC9Instance
      LabIdeInstanceProfileName:
        Ref: EksWorkshopC9InstanceProfile
      LabIdeInstanceProfileArn:
        Fn::GetAtt:
        - EksWorkshopC9InstanceProfile
        - Arn
      SsmDocument:
        Ref: EksWorkshopC9SSMDocument

  EksWorkshopC9BootstrapInstanceLambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.lambda_handler
      Role:
        Fn::GetAtt:
        - EksWorkshopC9LambdaExecutionRole
        - Arn
      Runtime: python3.9
      Environment:
        Variables:
          DiskSize:
            Ref: EksWorkshopC9InstanceVolumeSize
      MemorySize: 256
      Timeout: '900'
      Code:
        ZipFile: |
          from __future__ import print_function
          import boto3
          import json
          import os
          import time
          import traceback
          import cfnresponse
          import logging
          logger = logging.getLogger(__name__)
          
          def lambda_handler(event, context):
              print(event.values())
              print('context: {}'.format(context))
              responseData = {}

              status = cfnresponse.SUCCESS
              
              if event['RequestType'] == 'Delete':
                  responseData = {'Success': 'Custom Resource removed'}
                  cfnresponse.send(event, context, status, responseData, 'CustomResourcePhysicalID')              
              else:
                  try:
                      # Open AWS clients
                      ec2 = boto3.client('ec2')
                      ssm = boto3.client('ssm')
          
                      # Get the InstanceId of the Cloud9 IDE
                      instance = ec2.describe_instances(Filters=[{'Name': 'tag:Name','Values': ['aws-cloud9-'+event['ResourceProperties']['Cloud9Name']+'-'+event['ResourceProperties']['EnvironmentId']]}])['Reservations'][0]['Instances'][0]
                      print('instance: {}'.format(instance))
                      instance_id = instance['InstanceId']
          
                      # Create the IamInstanceProfile request object
                      iam_instance_profile = {
                          'Arn': event['ResourceProperties']['LabIdeInstanceProfileArn'],
                          'Name': event['ResourceProperties']['LabIdeInstanceProfileName']
                      }
                      print('Found IAM instance profile: {}'.format(iam_instance_profile))

                      time.sleep(10)

                      print('Waiting for the instance to be ready...')
          
                      # Wait for Instance to become ready before adding Role
                      instance_state = instance['State']['Name']
                      print('instance_state: {}'.format(instance_state))
                      while instance_state != 'running':
                          time.sleep(5)
                          instance_state = ec2.describe_instances(InstanceIds=[instance_id])
                          print('instance_state: {}'.format(instance_state))

                      print('Instance is ready')

                      associations = ec2.describe_iam_instance_profile_associations(
                          Filters=[
                              {
                                  'Name': 'instance-id',
                                  'Values': [instance_id],
                              },
                          ],
                      )

                      if len(associations['IamInstanceProfileAssociations']) > 0:
                        print('Replacing existing IAM profile...')
                        for association in associations['IamInstanceProfileAssociations']:
                          if association['State'] == 'associated':
                            print("{} is active with state {}".format(association['AssociationId'], association['State']))
                            ec2.replace_iam_instance_profile_association(AssociationId=association['AssociationId'], IamInstanceProfile=iam_instance_profile)
                      else:
                        print('Associating IAM profile...')
                        ec2.associate_iam_instance_profile(IamInstanceProfile=iam_instance_profile, InstanceId=instance_id)

                      block_volume_id = instance['BlockDeviceMappings'][0]['Ebs']['VolumeId']

                      block_device = ec2.describe_volumes(VolumeIds=[block_volume_id])['Volumes'][0]

                      DiskSize = int(os.environ['DiskSize'])
                      if block_device['Size'] < DiskSize:
                        ec2.modify_volume(VolumeId=block_volume_id,Size=DiskSize)
                        print('Modifying block volume: {}'.format(block_volume_id))
                        time.sleep(10)

                        for i in range(1, 30):
                          response = ec2.describe_volumes_modifications(
                            VolumeIds=[
                              block_volume_id
                            ]
                          )
                          modify_state = response['VolumesModifications'][0]['ModificationState']
                          if modify_state != 'modifying':
                            print('Volume has been resized')
                            break
                          time.sleep(10)
                      else:
                        print('Volume is already sized')

                      # Reboot is required to avoid weird race condition with IAM role and SSM agent
                      # It also causes the file system to expand in the OS
                      print('Rebooting instance')

                      ec2.reboot_instances(
                          InstanceIds=[
                              instance_id,
                          ],
                      )

                      time.sleep(60)

                      print('Waiting for instance to come online in SSM...')

                      for i in range(1, 60):
                        response = ssm.describe_instance_information(Filters=[{'Key': 'InstanceIds', 'Values': [instance_id]}])
                        if len(response["InstanceInformationList"]) == 0:
                          print('No instances in SSM')
                        elif len(response["InstanceInformationList"]) > 0 and \
                              response["InstanceInformationList"][0]["PingStatus"] == "Online" and \
                              response["InstanceInformationList"][0]["InstanceId"] == instance_id:
                          print('Instance is online in SSM')
                          break
                        time.sleep(10)

                      ssm_document = event['ResourceProperties']['SsmDocument']

                      print('Sending SSM command...')

                      response = ssm.send_command(
                        InstanceIds=[instance_id],
                        DocumentName=ssm_document)

                      command_id = response['Command']['CommandId']

                      waiter = ssm.get_waiter('command_executed')

                      waiter.wait(
                        CommandId=command_id,
                        InstanceId=instance_id,
                        WaiterConfig={
                          'Delay': 10,
                          'MaxAttempts': 30
                        }
                      )
  
                      responseData = {'Success': 'Started bootstrapping for instance: '+instance_id}
                      cfnresponse.send(event, context, status, responseData, 'CustomResourcePhysicalID')
                      
                  except Exception as e:
                      status = cfnresponse.FAILED
                      print(traceback.format_exc())
                      responseData = {'Error': traceback.format_exc(e)}
                  finally:
                      cfnresponse.send(event, context, status, responseData, 'CustomResourcePhysicalID')

  EksWorkshopC9SSMDocument: 
    Type: AWS::SSM::Document
    Properties: 
      DocumentType: Command
      DocumentFormat: YAML
      Content: 
        schemaVersion: '2.2'
        description: Bootstrap Cloud9 Instance
        mainSteps:
        - action: aws:runShellScript
          name: EksWorkshopC9bootstrap
          inputs:
            runCommand:
            - !Sub |
                set -e
      
                export AWS_REGION="${AWS::Region}"
                export REPOSITORY_OWNER="${RepositoryOwner}"
                export REPOSITORY_NAME="${RepositoryName}"
                export REPOSITORY_REF="${RepositoryRef}"
                export CLOUD9_ENVIRONMENT_ID="${EksWorkshopC9Instance}"
                export RESOURCES_PRECREATED="${ResourcesPrecreated}"

                curl -fsSL https://raw.githubusercontent.com/${RepositoryOwner}/${RepositoryName}/${RepositoryRef}/lab/scripts/installer.sh | bash

                sudo -E -H -u ec2-user bash -c "curl -fsSL https://raw.githubusercontent.com/${RepositoryOwner}/${RepositoryName}/${RepositoryRef}/lab/scripts/setup.sh | bash"

  EksWorkshopC9InstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: "/"
      Roles:
      - Ref: EksWorkshopC9Role

  EksWorkshopC9Instance:
    Type: AWS::Cloud9::EnvironmentEC2
    Properties:
      Description: AWS Cloud9 instance for EKS Workshop
      ImageId: amazonlinux-2-x86_64
      AutomaticStopTimeMinutes: 3600
      InstanceType:
        Ref: EksWorkshopC9InstanceType
      SubnetId: !If [ IsCloud9SubnetSpecified, !Ref Cloud9Subnet, !Ref "AWS::NoValue" ]
      Name: !If [ IsCloud9NotNamed, !Ref AWS::StackName, !Ref Cloud9Name ]
      OwnerArn: !If [ Create3rdPartyResources, !Ref WorkshopOwnerArn, !Ref "AWS::NoValue" ]
      Tags: 
        - Key: SSMBootstrap
          Value: Active

Outputs:
  Cloud9RoleArn:
    Description: The ARN of the IAM role assigned to the Cloud9 instance
    Value: !GetAtt EksWorkshopC9Role.Arn
    
  Cloud9InstanceName:
    Description: Name of the Cloud9 EC2 instance
    Value: !Sub 'aws-cloud9-${EksWorkshopC9Instance.Name}-${EksWorkshopC9Instance}'

  Cloud9Url:
    Value: !Sub 'https://${AWS::Region}.console.aws.amazon.com/cloud9/ide/${EksWorkshopC9Instance}?region=${AWS::Region}'