AWSTemplateFormatVersion: '2010-09-09'
Description: The template creates the Spoke VPC for the squid proxy using the AWS Transit Gateway and Squid.
Parameters:
  CidrBlock:
    Default: '10.10.0.0/16'
    Description: The CIDR Block for the VPC
    Type: String
  SubnetCidrBlock:
    Default: '10.10.0.0/24'
    Description: The CIDR Block for the Subnet
    Type: String
  SubnetCidrBlock2:
    Default: '10.10.1.0/24'
    Description: The CIDR Block for the Subnet
    Type: String
  ParentStackName:
    Description: The Stack that created the Transit Gateway
    Type: String
  AmazonLinuxAMI:
    Type: 'AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>'
    Default: '/aws/service/ami-amazon-linux-latest/amzn2-ami-hvm-x86_64-gp2'
    Description: DO NOT CHANGE - AMI Used for test instance.

Conditions: 
  TwoAZs:  !Not [!Equals [ !Ref SubnetCidrBlock2, '' ]]

Resources:

  rInstanceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - ec2.amazonaws.com
          Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonEC2RoleforSSM
      Path: "/"

  rInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: "/"
      Roles:
      - !Ref rInstanceRole

  myEC2TestInstance:
    Type: AWS::EC2::Instance
    Properties:
      IamInstanceProfile: !Ref rInstanceProfile
      ImageId: !Ref AmazonLinuxAMI
      InstanceType: t2.micro
      Monitoring: false
      SubnetId: !Ref ProtectedSubnet
      UserData: !Base64 |
        #!/bin/bash -ex
        yum -y install links
      Tags:
        - Key: Name
          Value: SpokeTestInstance

  VPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: !Ref CidrBlock
      Tags:
        - Key: Application
          Value: !Ref 'AWS::StackId'
        - Key: Name
          Value: !Sub Spoke-${AWS::StackName}
  ProtectedSubnet:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref 'VPC'
      CidrBlock: !Ref SubnetCidrBlock
      AvailabilityZone: 
        Fn::Select: 
          - 0
          - Fn::GetAZs: ""
      Tags:
        - Key: Application
          Value: !Ref 'AWS::StackId'
        - Key: Name
          Value: !Sub Protected-AZ1-${AWS::StackName}
  ProtectedSubnet2:
    Type: AWS::EC2::Subnet
    Condition: TwoAZs
    Properties:
      VpcId: !Ref 'VPC'
      CidrBlock: !Ref SubnetCidrBlock2
      AvailabilityZone: 
        Fn::Select: 
          - 1
          - Fn::GetAZs: ""
      Tags:
        - Key: Application
          Value: !Ref 'AWS::StackId'
        - Key: Name
          Value: !Sub Protected-AZ2-${AWS::StackName}
  ProtectedRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref 'VPC'
      Tags:
        - Key: Application
          Value: !Ref 'AWS::StackId'
        - Key: Name
          Value: !Sub Protected-${AWS::StackName}
  SubnetRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref 'ProtectedSubnet'
      RouteTableId: !Ref 'ProtectedRouteTable'
  Subnet2RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Condition: TwoAZs
    Properties:
      SubnetId: !Ref 'ProtectedSubnet2'
      RouteTableId: !Ref 'ProtectedRouteTable'
  NetworkAcl:
    Type: AWS::EC2::NetworkAcl
    Properties:
      VpcId: !Ref 'VPC'
      Tags:
        - Key: Application
          Value: !Ref 'AWS::StackId'
  InboundHTTPNetworkAclEntry:
    Type: AWS::EC2::NetworkAclEntry
    Properties:
      NetworkAclId: !Ref 'NetworkAcl'
      RuleNumber: '100'
      Protocol: '-1'
      RuleAction: allow
      Egress: 'false'
      CidrBlock: 0.0.0.0/0
  OutboundHTTPNetworkAclEntry:
    Type: AWS::EC2::NetworkAclEntry
    Properties:
      NetworkAclId: !Ref 'NetworkAcl'
      RuleNumber: '100'
      Protocol: '-1'
      RuleAction: allow
      Egress: 'true'
      CidrBlock: 0.0.0.0/0
  PublicEgressVpcSubnetNetworkAclAssociation:
    Type: AWS::EC2::SubnetNetworkAclAssociation
    Properties:
      SubnetId: !Ref 'ProtectedSubnet'
      NetworkAclId: !Ref 'NetworkAcl'
  PublicSecurity2VpcSubnetNetworkAclAssociation:
    Type: AWS::EC2::SubnetNetworkAclAssociation
    Condition: TwoAZs
    Properties:
      SubnetId: !Ref 'ProtectedSubnet2'
      NetworkAclId: !Ref 'NetworkAcl'
  SpokeAttachment:
    Type: "AWS::EC2::TransitGatewayAttachment"
    Properties:
      SubnetIds: 
       !If 
        - TwoAZs
        -
          - !Ref 'ProtectedSubnet'
          - !Ref 'ProtectedSubnet2'
        -
          - !Ref 'ProtectedSubnet'
      Tags: 
        - Key: Name
          Value: !Sub Spoke-${AWS::StackName}
      TransitGatewayId: 
        Fn::ImportValue:
          !Sub "${ParentStackName}-TransitGateway"
      VpcId: !Ref VPC
  SpokeTransitGatewayRoutePropagation:
    Type: "AWS::EC2::TransitGatewayRouteTablePropagation"
    Properties:
      TransitGatewayAttachmentId: !Ref SpokeAttachment
      TransitGatewayRouteTableId:
        Fn::ImportValue:
          !Sub "${ParentStackName}-TgwRouteTable"
  EgressVpcTransitGatewayRoutePropagation:
    Type: "AWS::EC2::TransitGatewayRouteTablePropagation"
    Properties:
      TransitGatewayAttachmentId: !Ref SpokeAttachment
      TransitGatewayRouteTableId:
        Fn::ImportValue:
          !Sub "${ParentStackName}-TgwEgressVpcRouteTable"
  SpokeVpcTgwAssociation:
    Type: "AWS::EC2::TransitGatewayRouteTableAssociation"
    Properties:
      TransitGatewayAttachmentId: !Ref SpokeAttachment
      TransitGatewayRouteTableId:
        Fn::ImportValue:
          !Sub "${ParentStackName}-TgwRouteTable"

  UpdateTgwRoutesFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonEC2FullAccess
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
  UpdateTgwRoutesFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      Handler: index.process_cfn
      Runtime: python2.7
      Description: ''
      MemorySize: 512
      Timeout: 240
      Role: !GetAtt UpdateTgwRoutesFunctionRole.Arn
      Code:
        ZipFile: |
          import boto3
          import random
          import string
          import uuid
          import httplib
          import urlparse
          import json
          import base64
          import hashlib
          import os
          import cfnresponse

          ec2 = boto3.resource('ec2')
          ec2client = boto3.client('ec2')

          def lambda_handler(event, context):
              try:
                  return process_cfn(event, context)
              except Exception as e:
                  print("EXCEPTION", e)
                  send_response(event, {
                      'StackId': event['StackId'],
                      'RequestId': event['RequestId'],
                      'LogicalResourceId': event['LogicalResourceId']
                      }, "FAILED")

          def process_cfn(event, context):
              print("Received event: " + json.dumps(event, indent=2))

              tgw = event['ResourceProperties']['TransitGateway']
              spokeVpcObj = ec2.Vpc(event['ResourceProperties']['SpokeVpc'])
              EgressVpcObj = ec2.Vpc(event['ResourceProperties']['EgressVpc'])
              rts = []
              rts.append(event['ResourceProperties']['PrivateRouteTable'])
              rts.append(event['ResourceProperties']['PrivateRouteTable2'])
              rts.append(event['ResourceProperties']['TGWAttchRouteTable'])
              rts.append(event['ResourceProperties']['TGWAttchRouteTable2'])
            
              spokeRt = event['ResourceProperties']['SpokeProtectedRouteTable']
              
              cidr = spokeVpcObj.cidr_block
              if event['RequestType'] == 'Delete':
                  for rt in rts:
                      if rt != "SingleAZ":
                          ec2client.delete_route(
                              DestinationCidrBlock= cidr,
                              RouteTableId=rt
                          )

                  delete_all_routes(spokeRt)
                  responseData = {}
                  responseData['Data'] = "Delete Succcessful"
                  return cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData)

              for rt in rts:
                  if rt != "SingleAZ":
                      ec2.RouteTable(rt).create_route(
                          DestinationCidrBlock=cidr,
                          GatewayId=tgw
                      )
              
              
              spoke_route_table = ec2.RouteTable(spokeRt)
              cidr = EgressVpcObj.cidr_block

              route = spoke_route_table.create_route(
                  DestinationCidrBlock=cidr,
                  GatewayId=tgw
              )
              
              route = spoke_route_table.create_route(
                  DestinationCidrBlock='0.0.0.0/0',
                  GatewayId=tgw
              )
              
              responseData = {}
              responseData['Data'] = "Create Succcessful"
              return cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData)

          def delete_all_routes(routeTableId):
              source_route_table = ec2.RouteTable(routeTableId)
              
              for route in source_route_table.routes:

                  if route.origin == "CreateRoute":
                      response = ec2client.delete_route(
                          DestinationCidrBlock= route.destination_cidr_block,
                          RouteTableId=routeTableId
                      )
              return ""
  UpdateTgwRoutes:
    Type: Custom::UpdateTgwRoutes
    DependsOn:
      - UpdateTgwRoutesFunction
      - SpokeTransitGatewayRoutePropagation
      - SpokeVpcTgwAssociation
    Properties: 
      ServiceToken: !GetAtt UpdateTgwRoutesFunction.Arn
      EgressVpc:
        Fn::ImportValue:
          !Sub "${ParentStackName}-EgressVpc"
      SpokeVpc: !Ref VPC
      PrivateRouteTable: 
        Fn::ImportValue:
          !Sub "${ParentStackName}-PrivateRouteTable"
      TGWAttchRouteTable:
        Fn::ImportValue:
          !Sub "${ParentStackName}-TGWAttchRouteTable"
      PrivateRouteTable2: 
        Fn::ImportValue:
          !Sub "${ParentStackName}-PrivateRouteTable2"
      TGWAttchRouteTable2: 
        Fn::ImportValue:
          !Sub "${ParentStackName}-TGWAttchRouteTable2"
      SpokeProtectedRouteTable: !Ref ProtectedRouteTable
      TransitGateway:
        Fn::ImportValue:
          !Sub "${ParentStackName}-TransitGateway"

Outputs:

  SpokeProtectedRouteTable:
    Value: !Ref ProtectedRouteTable
  SpokeAttachment:
    Value: !Ref SpokeAttachment
    Description: The spoke attachment to the TransitGateway
  myEC2TestInstance:
      Value: !Ref myEC2TestInstance
      Description: Test Instance for internet access.