zs_cc_cf_template_zscc_macro.yaml

AWSTemplateFormatVersion: 2010-09-09
Description: Zscaler Cloud Connector Pre-Deployment Macro Template
Metadata:
  LICENSE: 'Apache License, Version 2.0'
  cfn-lint:
    config:
      ignore_checks:
        - E1029 # Fn::Sub check for variable in string
Resources:
  ZSCCMacroExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      Path: /
      Policies:
        - PolicyName: CustomLambdaEC2DescribePolicy
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                Resource:
                  - !Sub 'arn:aws:logs:*:${AWS::AccountId}:log-group:/aws/lambda/zscc-macro'
              - Effect: Allow
                Action:
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource:
                  - !Sub 'arn:aws:logs:*:${AWS::AccountId}:log-group:*'
              - Effect: Allow
                Action:
                  - ec2:AssignPrivateIpAddresses
                  - ec2:CreateNetworkInterface
                  - ec2:DeleteNetworkInterface
                  - ec2:DescribeInstances
                  - ec2:DescribeNetworkInterfaces
                  - ec2:DescribeSubnets
                  - ec2:DescribeInstanceAttribute
                  - ec2:DescribeInstanceTypes
                  - ec2:DescribeInstanceStatus
                  - ec2:DescribeRouteTables
                  - ec2:ReplaceRoute
                  - ec2:UnassignPrivateIpAddresses
                  - ec2:DescribeImages
                  - lambda:InvokeFunction
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                  - secretsmanager:ListSecrets
                  - secretsmanager:GetSecretValue
                Resource: "*"
              - Effect: Allow
                Action:
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                Resource: 'arn:aws:logs:*:*:*'
  
  ZSCCMacroFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: zscc-macro
      Handler: index.lambda_handler
      Description: Retrieves EC2 Service ENI Properties
      Timeout: 30
      Role: !GetAtt
        - ZSCCMacroExecutionRole
        - Arn
      Runtime: python3.7
      Code:
        ZipFile: |
          import json
          import logging
          import boto3
          import copy
          import collections
          import re
          import copy
          import pprint

          logger = logging.getLogger()
          logger.setLevel(logging.INFO)

          # ONE TIME INITS
          EC2_RESOURCE = boto3.resource('ec2')
          EC2_CLIENT = boto3.client('ec2')
          CCINSTANCETYPE2INTS = {1: "Small", 2: "Medium", 3: "Large"}
          IPV4PATTERN = re.compile("^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])$")

          def retrieve_eni_properties(event):
              try:
                responseData = {"status": "success", "message": 'Retrieved ENI Properties!'}
                VpcIds = []
                SubnetIds = []
                Targets = []
                InstanceTypes = []
                InstanceNumSrvcInts = []
                SubnetsByAZ = collections.defaultdict(list)
                instances = event.get("templateParameterValues").get("Ec2Instances")
                for instance_id in instances:
                  instance = EC2_RESOURCE.Instance(instance_id)
                  instance_type = instance.instance_type
                  logger.info(f"Instance: {instance_id}, Type: {instance_type}")
                  if instance_type not in InstanceTypes:
                    InstanceTypes.append(instance_type)
                  # Only works for static service NIC at index 0
                  enis = [e for e in
                    instance.network_interfaces_attribute if e.get("Attachment").get("DeviceIndex") == 0]
                  # Let's keep track of how many enis this instance has
                  if len(enis) not in InstanceNumSrvcInts:
                    InstanceNumSrvcInts.append(len(enis))
                  if len(enis):
                    for eni in enis:
                      ip = [i.get("PrivateIpAddress") for i in 
                        eni.get("PrivateIpAddresses") if i.get("Primary") == True][0]
                      Targets.append({"Id": ip})
                      subnetId = eni.get("SubnetId")
                      subnet = EC2_RESOURCE.Subnet(subnetId)
                      if subnetId not in SubnetIds:
                        SubnetIds.append(subnetId)
                        SubnetsByAZ[re.sub("\-", "", subnet.availability_zone)].append(subnetId)
                      vpc = eni.get("VpcId")
                      if vpc not in VpcIds:
                        VpcIds.append(vpc)
                  else:
                    err = f"Instance {instance} does not have a multiple enis"
                    logger.critical(err)
                    responseData["status"] = "failed"
                    responseData["errorMessage"] = err  
                    break
                responseData["Targets"] = Targets
                responseData["SubnetIds"]= SubnetIds
                responseData["SubnetsByAZ"] = SubnetsByAZ
                if len(InstanceTypes) > 1:
                  err = f"All Cloud Connectors should be of the same instance type. Multiple found {InstanceTypes}"
                  logger.critical(err)
                  responseData["status"] = "failed"
                  responseData["errorMessage"] = err
                if len(InstanceNumSrvcInts) > 1:
                  err = f"All Cloud Connectors should be of the same CC Instance Type. Multiple found {[CCINSTANCETYPE2INTS.get(i) for i in InstanceNumSrvcInts]}"
                  logger.critical(err)
                  responseData["status"] = "failed"
                  responseData["errorMessage"] = err
                if len(VpcIds) > 1:
                  err = f"All Cloud Connectors should be in the same VPC, Multiple found {VpcIds}"
                  logger.critical(err)
                  responseData["status"] = "failed"
                  responseData["errorMessage"] = err
                elif len(VpcIds) == 1:
                  responseData["VPCID"] = VpcIds[0]
                if responseData.get('status') == "success":
                  logger.info('Retrieved ENI Properties!')
                return responseData
              except Exception as e:
                logger.critical(f"Error Occured: {e}")

          def add_required_gwlb_resources(template, responseData):
              res = {
                "status": "success",
                "message": "Template processed successfully"
              }
              try:
                resources = template.get("Resources")
                outputs = template.get("Outputs")
                # Resource Keys to touch are
                # ZSCCGWLBTargetGroup, ZSCCVPCEP, ZSCCGWLB
                ZSCCGWLBTargetGroup = resources.get("ZSCCGWLBTargetGroup")
                ZSCCGWLBTargetGroup.get("Properties")["Targets"] = responseData.get("Targets")
                ZSCCGWLBTargetGroup.get("Properties")["VpcId"] = responseData.get("VPCID")
                ZSCCGWLB = resources.get("ZSCCGWLB")
                ZSCCGWLB.get("Properties")["Subnets"] = responseData.get("SubnetIds")
                ZSCCVPCEPTemplate = resources.pop("ZSCCVPCEP")
                ZSCCVPCEPOutputTemplate = outputs.pop("ZSCCVPCEP")
                ZSCCVPCEPTemplate.get("Properties")["VpcId"] = responseData.get("VPCID")
                for az, subnets in responseData.get("SubnetsByAZ").items():
                  ep_name = f"ZSCCVPCEP{az}"
                  ep_resource = copy.deepcopy(ZSCCVPCEPTemplate)
                  ep_resource.get("Properties")["SubnetIds"] = subnets[:1]
                  resources[ep_name] = ep_resource
                  ep_output = copy.deepcopy(ZSCCVPCEPOutputTemplate)
                  ep_output["Description"] = f"{ep_output['Description']} in {az}"
                  ep_output["Value"] = {"Ref": ep_name}
                  outputs[ep_name] = ep_output
                res["fragment"] = template
              except Exception as e:
                logger.critical(f"Error Occured: {e}")
                res["status"] = "failed"
                res["errorMessage"] = e
              return res

          def zsccgwlb(event, context):
              resp = {
                  "requestId": event.get("requestId")
                }
              try:
                responseData = retrieve_eni_properties(event)
                resp.update(responseData)
                if resp.get("status") == "success":
                  template = event.get("fragment")
                  res = add_required_gwlb_resources(template=template, responseData=responseData)
                  resp.update(res)
              except Exception as e:
                logger.critical(e)
              return resp

          def retrive_ami_id_and_secret_arn(event):

              responseData = {"status": "success", "message": 'Retrieved AMI ID!'}

              # get the ImageCode from the template
              try:
                
                if event.get("templateParameterValues").get("ZscalerOsAmi"):
                  ImageId = event.get("templateParameterValues").get("ZscalerOsAmi")
                else:
                  ImageCode = event["fragment"]["Mappings"]["Product2Code"]["CloudConnector"]["Code"]
                  logger.debug(f"Image Code: {ImageCode}")
                  # use the ImageCode to get ami id
                  filters = [{"Name": "product-code", "Values":[ImageCode]}]
                  response = sorted(EC2_CLIENT.describe_images(Owners=["aws-marketplace"], Filters=filters).get("Images"), key=lambda x:  x.get("CreationDate"))
                  ImageId = response[-1].get("ImageId")
                
              except Exception as e:
                err = f"Retrieve AMI ID errored: {e}"
                logger.info(err)
                responseData = {"status": "failed", "errorMessage": f'{e}'}
                return responseData

              # Retrieve secrets manager ARN
              try:
                SecretManagerSecretName = event["templateParameterValues"].get("ZscalerSecretManagerSecretName", "ZS/CC/credentials")
                secretsmanager_client = boto3.client('secretsmanager')
                secret = secretsmanager_client.get_secret_value(SecretId=SecretManagerSecretName)
                
                responseData["ImageId"] = ImageId
                responseData["SecretArn"] = secret.get("ARN")
                
                logger.info('Retrieved AMI {0}'.format(responseData.get('ImageId')))
                logger.info('Retrieved Secret ARN {0}'.format(responseData.get('SecretArn')))
                logger.info(f'Sending response {responseData}')
                
                return responseData

              except Exception as e:
                err = f"Retriving Secret Manager ARN failed: {e}"
                logger.error(err)
                responseData = {"status": "failed", "errorMessage": f'{err}'}
                return responseData

          def add_required_simple_resources(template, responseData):
            res = {
                "status": "success",
                "message": "Simple Template processed successfully"
              }
              
            try:  
              # get the resources and outputs from the template
              resources = template.get("Resources")
              outputs = template.get("Outputs")

              # add SecretArn to CcIamRole Policy
              # first fetch the policy document
              policy_doc = resources["CcGetSecretsPolicy"]["Properties"]["PolicyDocument"]
              
              # update the document
              policy_doc["Statement"][0]["Resource"] = responseData["SecretArn"]
              resources["CcGetSecretsPolicy"]["Properties"]["PolicyDocument"] = policy_doc

              # add ami id to ec-2 resource
              resources["Ec2Instance"]["Properties"]["ImageId"] = responseData["ImageId"]
              
              # update template with new resources
              template["Resources"] = resources

              # add ami id in use to outputs
              outputs["ImageIdInUse"]["Value"] = responseData["ImageId"]

              # update outputs in template
              template["Outputs"] = outputs

              # update response 
              res["fragment"] = template

              return res

            except Exception as e:
              res = {
                "status": "failure",
                "errorMessage": f"{e}"
              }
              return res

          def zsccsimple(event, context):
              resp = {
                  "requestId": event.get("requestId")
                }
              
              try:
                responseData = retrive_ami_id_and_secret_arn(event)
                resp.update(responseData)
                if resp.get("status") == "success":
                  template = event.get("fragment")
                  res = add_required_simple_resources(template=template, responseData=responseData)
                  resp.update(res)
              
              except Exception as e:
                logger.critical(e)
              return resp

          def get_secondary_ip_from_service_eni(event):

            responseData = {"status": "success", "message": 'Retrieved ENI properties!'}

            try:
              # get ENI from template
              serviceENI = event['templateParameterValues'].get('ServiceENI')
              logger.info(f'ENI as per template is {serviceENI}')
              
              # use ENI to get secondary service IP addrs              
              client = boto3.client('ec2')
              response = client.describe_network_interfaces(NetworkInterfaceIds=[serviceENI])

              # update responseData
              responseData['TargetIp']= response.get('NetworkInterfaces')[0].get('PrivateIpAddresses')[1]['PrivateIpAddress']
              responseData['VpcId']= response.get('NetworkInterfaces')[0].get('VpcId')
              responseData['SubnetId']= response.get('NetworkInterfaces')[0].get('SubnetId')
              logger.info('Retrieved ENI Properties!')

            except Exception as err:
              logger.error(f'function is -> get_secondary_ip_from_service_eni, error is {err}')
              responseData = {"status": "failure", "errorMessage": str(err)}
            
            return responseData

          def add_required_zpa_resources(template, vpcId, ResolverEPSubnets, TargetIps, app_domains_list):
            
            res = {
                "status": "success",
                "message": "ZPA Template processed successfully"
              }

            # get current resources and add to it later 
            resources = template.get("Resources")

            # outline of domain rule to add 
            domain_rule_resource_outline = resources.get("AppSegmentResolverDomainRule")

            # outline of domain rule association to add 
            domain_rule_assn_resource_outline = resources.get("AppSegmentResolverDomainRuleAssociation")

            try:
              # add extra app domain rules and associations

              for i in range(len(app_domains_list)):

                app_domain = app_domains_list[i]
                tag_name = "${AWS::StackName}-"+f"ZSCCAppSegment{i}"
                
                rule_resource_name = f"AppSegmentResolverDomainRule{i}"
                association_resource_name = f"AppSegmentResolverDomainRule{i}Association"
                
                # REF! https://stackoverflow.com/questions/3975376/understanding-dict-copy-shallow-or-deep
                # new rule res
                new_rule_res = copy.deepcopy(domain_rule_resource_outline)
                
                # modify fields
                new_rule_res["Properties"]["DomainName"] = app_domain
                new_rule_res["Properties"]["Name"] = {"Fn::Sub": tag_name}
                new_rule_res["Properties"]["Tags"][0]["Value"] = {"Fn::Sub": tag_name}
                new_rule_res["Properties"]["TargetIps"] = list()
                for targetIp in TargetIps:
                  new_rule_res["Properties"]["TargetIps"].append({"Ip": targetIp, "Port": 53})

                # add new rule resource
                resources[rule_resource_name] = new_rule_res

                # new association to rule
                new_assn_res = copy.deepcopy(domain_rule_assn_resource_outline)

                # modify fields
                new_assn_res["Properties"]["Name"] = {"Fn::Sub":tag_name+"Association"}
                new_assn_res["Properties"]["ResolverRuleId"]["Ref"] = rule_resource_name
                new_assn_res["Properties"]["VPCId"] = vpcId

                # add new assn resource
                resources[association_resource_name] = new_assn_res

              # modify existing resources
              resources["AppSegmentResolverSecurityGroup"]["Properties"]["VpcId"] = vpcId
              
              resources["AppSegmentOutboundResolver"]["Properties"]["IpAddresses"][0] = {"SubnetId" : ResolverEPSubnets[0]}
              resources["AppSegmentOutboundResolver"]["Properties"]["IpAddresses"][1] = {"SubnetId" : ResolverEPSubnets[len(ResolverEPSubnets) > 1]}
              
              for subnetId in ResolverEPSubnets[2:]:
                resources["AppSegmentOutboundResolver"]["Properties"]["IpAddresses"].append({"SubnetId" : subnetId})

              resources["AppSegmentResolverZSCloudDomainRuleAssociation"]["Properties"]["VPCId"] = vpcId
              resources["AppSegmentResolverZPACloudDomainRuleAssociation"]["Properties"]["VPCId"] = vpcId
              resources["AppSegmentResolverZScalerComDomainRuleAssociation"]["Properties"]["VPCId"] = vpcId
              resources["AppSegmentResolverFreeBSDRuleAssociation"]["Properties"]["VPCId"] = vpcId
              resources["AppSegmentResolverNtpOrgDomainRuleAssociation"]["Properties"]["VPCId"] = vpcId
              resources["AppSegmentResolverAmazonAWSRuleAssociation"]["Properties"]["VPCId"] = vpcId

              # pop outline resources
              resources.pop("AppSegmentResolverDomainRule")
              resources.pop("AppSegmentResolverDomainRuleAssociation")

              # modify template and return modified template
              template["Resources"] = resources

              res["fragment"] = template


            except Exception as e:
              logger.error(f'function is -> add_required_zpa_resources, error is {e}')
              res["status"] = "failure"
              res["errorMessage"] = e
            
            return res

          def zscczpa(event, context):
            
            logger.info('got event {}'.format(event))
            resp = {
                  "requestId": event.get("requestId")
                }
            
            # get secondary IP from service ENI
            try :
              # responseData = get_secondary_ip_from_service_eni(event)
              ResolverEPSubnets = event.get("templateParameterValues").get("ResolverEPSubnets")
              logger.info(f"ResolverEPSubnets {ResolverEPSubnets}")
              if "" in ResolverEPSubnets or len(ResolverEPSubnets) < 1 or len(ResolverEPSubnets) > 6:
                msg = f"ResolverEPSubnets {ResolverEPSubnets} count should be between 1 and 6 inclusive."
                resp["status"] = "failed"
                resp["errorMessage"] = msg
                return resp
              
              TargetIps = event.get("templateParameterValues").get("TargetIPs")
              logger.info(f"TargetIps {TargetIps}")
              if type(TargetIps) is not list or len(TargetIps) < 1 or len(TargetIps) > 6:
                msg = f"TargetIps {TargetIps} count should be between 1 and 6 inclusive."
                resp["status"] = "failed"
                resp["errorMessage"] = msg
                return resp
              elif None in [IPV4PATTERN.fullmatch(ip) for ip in TargetIps]:
                msg = f"All TargetIps {TargetIps} should be valid ip addresses"
                resp["status"] = "failed"
                resp["errorMessage"] = msg
                return resp
              vpcId = event.get("templateParameterValues").get("VPCID")
              
              # parse input parameters to get all app domains
              app_domains_list = event["templateParameterValues"].get("ApplicationDomainNames")
              logger.info(app_domains_list)

              # update the template by dynamically adding resources
              template = event.get("fragment")
              res = add_required_zpa_resources(template, vpcId, ResolverEPSubnets, TargetIps, app_domains_list)

              resp.update(res)

            except Exception as e:
                logger.critical(e)
            
            return resp

          def is_valid_sg(sgid):
            ## Validate Security Groups
            SEC_GROUP_REGEX = r"^sg-[0-9a-f]+$"
            # Regex first
            if not re.match(SEC_GROUP_REGEX, sgid):
              return False
            # Now send it to validate, this is not possible if user doesnt have permissions
            # securityGroups = EC2_CLIENT.describe_security_groups(GroupIds=[sgid]).get("SecurityGroups")
            # if len(securityGroups) != 1 or securityGroups[0].get("GroupId") != sgid:
            #   return False
            return True
            
          # def is_valid_iam_role_arn(arn):
          #   ## Validate IAM Role, this is not possible if user doesnt have permissions
          #   iam_role = EC2_RESOURCE.describe_iam_roles()

          def validate_inputs_asg(event, responseData):
            parameters = event.get("templateParameterValues")
            template = event.get("fragment")
            resources = template.get("Resources")
            outputs = template.get("Outputs")
            
            responseData = {
              "status": "success", 
              "message": "",
              "fragment": template
            }
            
            # Subnets In VPC and AZs- From CcSubnetIds and VpcEpSubnetIds
            subnetIds = list()
            subnetsByAz = collections.defaultdict(list)
            vpc = parameters.get("NetworkVpcId")
            availability_zones = parameters.get("NetworkAvailabilityZones")
            ccSubnetIds = parameters.get("NetworkCcSubnetIds", list())
            vpcEpSubnetIds = parameters.get("GwlbVpcEpSubnetIds", list())
            
            # Describe these Subnets
            try:
              allSubnetIds = ccSubnetIds+vpcEpSubnetIds if len(vpcEpSubnetIds[0]) > 0 else ccSubnetIds
              ec2_subnets = EC2_CLIENT.describe_subnets(SubnetIds=allSubnetIds)
              for ec2_subnet in ec2_subnets.get("Subnets"):
                subnetId = ec2_subnet.get("SubnetId")
                availability_zone = ec2_subnet.get('AvailabilityZone')
                if ec2_subnet.get("VpcId") != vpc:
                  responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
                  responseData["message"] += f"\nSubnet {ec2_subnet.get('SubnetId')} is not part of Vpc {vpc}."
                if ec2_subnet.get("AvailabilityZone") not in availability_zones:
                  responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
                  responseData["message"] += f"\nSubnet {subnetId}'s availability zone {availability_zone} is not part of given zones {availability_zones}."
                if subnetId not in subnetIds:
                  subnetIds.append(subnetId)
                  # We only want to add to subnetsByAz if we dont have vpcEpSubnets given or if 
                  # the subnet is part of vpcEpSubnets
                  if not len(vpcEpSubnetIds) or vpcEpSubnetIds[0] == '' or subnetId in vpcEpSubnetIds:
                    subnetsByAz[re.sub("\-", "", availability_zone)].append(subnetId)
              responseData["subnetIds"]= subnetIds
              responseData["subnetsByAz"] = subnetsByAz
            except Exception as e:
              err = f"\nValidating subnets errored: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
              
            # smallCCInstanceSize
            try:
              ALLOWED_SIZES = {
                "small": ["t3.medium", "m5.large", "c5.large", "c5a.large"]
              }
              ccInstanceSize = parameters.get("ZscalerCcInstanceSize")
              awsEc2InstanceType = parameters.get("ZscalerAwsEc2InstanceType")
              if awsEc2InstanceType not in ALLOWED_SIZES.get(ccInstanceSize):
                responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
                responseData["message"] += f"\naws Ec2 Insance Type {awsEc2InstanceType} is not valid for CcInstanceSize {ccInstanceSize}, allowed values are {ALLOWED_SIZES.get('ccInstanceSize')}."
            except Exception as e:
              err = f"Validating instance sizes errored: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
              
            # MgmtIntfSecurityGroupIsValid
            try:
              mgmtIntfSecurityGroup = parameters.get("ZscalerCcMgmtIntfSecurityGroup")
              if mgmtIntfSecurityGroup != "" and not is_valid_sg(sgid=mgmtIntfSecurityGroup):
                  responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
                  responseData["message"] += f"\nMgmt Intf Security Group specified {mgmtIntfSecurityGroup} should be a single valid AWS security group Id in the given region, please remove extraneous spaces, commas etc."
            except Exception as e:
              err = f"Validating Mgmt Security Group errored: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
                
            # SrvcIntfSecurityGroupIsValid
            try:
              srvcIntfSecurityGroup = parameters.get("ZscalerCcSrvcIntfSecurityGroup")
              if srvcIntfSecurityGroup != "" and not is_valid_sg(sgid=srvcIntfSecurityGroup):
                  responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
                  responseData["message"] += f"\nMgmt Intf Security Group specified {srvcIntfSecurityGroup} should be a single valid AWS security group Id in the given region, please remove extraneous spaces, commas etc."
            except Exception as e:
              err = f"Validating Srvc Security Group errored: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
            
            # Each AZ specified has at least one subnet to put into the GWLB, VPC ep subnets
            # override the CC Subnets. One subnet in each AZ is picked
            try:
              for az in availability_zones:
                if len(subnetsByAz.get(re.sub("\-", "", az), [])) == 0:
                  responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
                  responseData["message"] += f"\nAvailability Zone {az} does not have any subnets specified in Vpc Endpoint Subnets if specified or in CloudConnector Subnets"
            except Exception as e:
              err = f"Validating Availability Zones errored: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
            
            # If given, check if IAM Role exists if possible
            ## We wont do this because it requires user's ability to list roles etc
            ## which has great potential to break the script
            
            # If given, validate VpcEpServicePrincipals exist?
            
            if responseData.get("status") == "success":
              responseData["message"] = "All Inputs Validated successfully!"
            else:
              responseData["errorMessage"] = responseData["message"]
              
            logger.info(f"Validating Input: {responseData.get('status')}, Message: {responseData.get('message')}, Error: {responseData.get('errorMessage')}")
            
            return responseData

          def retrieve_data_asg(event, responseData):
            parameters = event.get("templateParameterValues")
            template = event.get("fragment")
            resources = template.get("Resources")
            outputs = template.get("Outputs")
            
            responseData = {
              "status": "success", 
              "message": "",
              "fragment": template
            }
            try:
              ami_and_secret_data = retrive_ami_id_and_secret_arn(event)
              responseData.update(ami_and_secret_data)
            except Exception as e:
              err = f"Retriving data failed: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
            
            if responseData.get("status") == "success":
              responseData["message"] = "All Data Retrieved successfully!"
            else:
              responseData["errorMessage"] = responseData["message"]
              
            logger.info(f"Retrieve Data: {responseData.get('status')}, Message: {responseData.get('message')}, Error: {responseData.get('errorMessage')}")
            
            
            return responseData

          def add_required_asg_resources(event, responseData):
            parameters = event.get("templateParameterValues")
            template = event.get("fragment")
            resources = template.get("Resources")
            outputs = template.get("Outputs")
            
            subnetsByAz = responseData.get("subnetsByAz")
            logger.info(f"Subnets By AZ: {subnetsByAz}")
            uniqueCcSubnetIds = list(set(parameters.get("NetworkCcSubnetIds")))
            
            # Launch template Image Id
            try:
              resources["CcLaunchTemplate"]["Properties"]["LaunchTemplateData"]["ImageId"] = responseData["ImageId"]
              outputs["ImageIdInUse"]["Value"] = responseData["ImageId"]
            except Exception as e:
              err = f"Adding ImageId to use failed: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
            
            # add SecretArn to CcIamRole Policy
            try:
              # This only works if the secretmanager policy is the first policy
              # first fetch the policy document
              policy_doc = resources["CcGetSecretsPolicy"]["Properties"]["PolicyDocument"]
              policy_doc_lambda = resources["CcGetSecretsPolicyForLambda"]["Properties"]["PolicyDocument"]
              
              # update the document
              policy_doc["Statement"][0]["Resource"] = responseData["SecretArn"]
              resources["CcGetSecretsPolicy"]["Properties"]["PolicyDocument"] = policy_doc
              policy_doc_lambda["Statement"][0]["Resource"] = responseData["SecretArn"]
              resources["CcGetSecretsPolicyForLambda"]["Properties"]["PolicyDocument"] = policy_doc_lambda
            except Exception as e:
              err = f"Adding SecretManger  to use failed: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
            
            # Asg Vpc Zone Identifier
            try:
              CcAutoScalingGroup = resources.get("CcAutoScalingGroup")
              CcAutoScalingGroup.get("Properties")["VPCZoneIdentifier"] = uniqueCcSubnetIds
              # One asg per az
              CcAutoScalingGroupTemplate = resources.pop("CcAutoScalingGroup")
              CcAsgAutoScalingPolicyTemplate = resources.pop("CcAsgAutoScalingPolicy")
              CcAsgWarmPoolTemplate = resources.pop("CcAsgWarmPool")
              CcAsgLambdaFunction = resources.get("CcAsgLambdaFunction")
              CcAsgEventBridgeInstanceTerminateLifecycleActionRule = resources.get("CcAsgEventBridgeInstanceTerminateLifecycleActionRule")
              asg_names = list()
              for az, subnets in subnetsByAz.items():
                asg_name = f"CcAutoScalingGroupFor{az}"
                asg_names.append(f"${{{asg_name}}}")
                asg_policy_name = f"CcAutoScalingGroupScalingPolicyFor{az}"
                asg_warmpool_name = f"CcAutoScalingGroupWarmPoolFor{az}"
                asg_resource = copy.deepcopy(CcAutoScalingGroupTemplate)
                asg_policy_resource = copy.deepcopy(CcAsgAutoScalingPolicyTemplate)
                asg_warmpool_resource = copy.deepcopy(CcAsgWarmPoolTemplate)
                asg_resource["Properties"]["AutoScalingGroupName"]["Fn::Sub"] = asg_resource["Properties"]["AutoScalingGroupName"]["Fn::Sub"] + f"For{az}"
                asg_policy_resource["Properties"]["AutoScalingGroupName"]["Ref"] = asg_name
                asg_policy_resource["Properties"]["TargetTrackingConfiguration"]["CustomizedMetricSpecification"]["Dimensions"][0]["Value"]["Ref"] = asg_name
                asg_warmpool_resource["Properties"]["AutoScalingGroupName"]["Ref"] = asg_name
                for lch in asg_resource["Properties"]["LifecycleHookSpecificationList"]:
                  lch["LifecycleHookName"]["Fn::Sub"] = lch["LifecycleHookName"]["Fn::Sub"] + f"For{az}"
                resources[asg_name] = asg_resource
                resources[asg_policy_name] = asg_policy_resource
                resources[asg_warmpool_name] = asg_warmpool_resource
              CcAsgLambdaFunction["Properties"]["Environment"]["Variables"]["ASG_NAMES"]["Fn::Sub"] = json.dumps(asg_names)
              CcAsgEventBridgeInstanceTerminateLifecycleActionRuleFnSub = json.loads(CcAsgEventBridgeInstanceTerminateLifecycleActionRule["Properties"]["EventPattern"]["Fn::Sub"])
              CcAsgEventBridgeInstanceTerminateLifecycleActionRuleFnSub["detail"]["AutoScalingGroupName"] = asg_names
              CcAsgEventBridgeInstanceTerminateLifecycleActionRule["Properties"]["EventPattern"]["Fn::Sub"] = json.dumps(CcAsgEventBridgeInstanceTerminateLifecycleActionRuleFnSub)

            except Exception as e:
              err = f"Modifying AutoScaling Group Resource failed: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
            
            # Vpc Ep
            try:
                # Resource Keys to touch are
                # CcGwlbTargetGroup, CcVpcEp, CcGwlb
                CcGwlbTargetGroup = resources.get("CcGwlbTargetGroup")
                # CcGwlbTargetGroup.get("Properties")["Targets"] = responseData.get("Targets")
                CcGwlbTargetGroup.get("Properties")["VpcId"] = parameters.get("NetworkVpcId")
                CcGwlb = resources.get("CcGwlb")
                CcGwlb.get("Properties")["Subnets"] = uniqueCcSubnetIds
                CcVpcEpTemplate = resources.pop("CcVpcEp")
                CcVpcEpOutputTemplate = outputs.pop("ZscalerCcVpcEp")
                CcVpcEpTemplate.get("Properties")["VpcId"] = parameters.get("NetworkVpcId")
                for az, subnets in subnetsByAz.items():                  
                  ep_name = f"ZSCCVPCEP{az}"
                  ep_resource = copy.deepcopy(CcVpcEpTemplate)
                  ep_resource.get("Properties")["SubnetIds"] = subnets[:1]
                  resources[ep_name] = ep_resource
                  ep_output = copy.deepcopy(CcVpcEpOutputTemplate)
                  ep_output["Description"] = f"{ep_output['Description']} in {az}"
                  ep_output["Value"] = {"Ref": ep_name}
                  outputs[ep_name] = ep_output
            except Exception as e:
              err = f"Modifying Gwlb Resourced failed: {e}"
              logger.error(err)
              responseData["status"] = ["failed", "success"][responseData["status"]=="success" and False]
              responseData["message"] += err
            
            
            template["Resources"] = resources
            template["Outputs"] = outputs
            
            responseData["fragment"] = template
            
            
            if responseData.get("status") == "success":
              responseData["message"] = "All required asg resources added/updated successfully!"
            else:
              responseData["errorMessage"] = responseData["message"]
            
            return responseData

          def zsccasg(event, context):
              resp = {
                  "requestId": event.get("requestId")
                }
              template = event.get("fragment")   
              try:

                # validate inputs
                responseData = validate_inputs_asg(event=event, responseData=resp)
                resp.update(responseData)

                # retrieve data
                if resp.get("status") == "success":
                  responseData = retrieve_data_asg(event=event, responseData=resp)
                  resp.update(responseData)

                # add/update template
                if resp.get("status") == "success":
                  responseData = add_required_asg_resources(event=event, responseData=resp)
                  logger.info(responseData)
                  resp.update(responseData)
                

              except Exception as e:
                
                logger.critical(e)
              
              return resp

          def lambda_handler(event, context):
              logger.info(json.dumps(event))
              operation = OPER_MAP.get(event.get("params").get("operation").lower())
              resp = operation(event, context)
              #logger.info(pprint.pformat(resp, indent=4))
              return resp

          # MAP
          OPER_MAP = {
            "zsccgwlb": zsccgwlb,
            "zsccsimple": zsccsimple,
            "zscczpa": zscczpa,
            "zsccasg": zsccasg
          }


  ZSCCMacroFunctionPermissions:
    Type: AWS::Lambda::Permission
    Properties:
      Action: 'lambda:InvokeFunction'
      FunctionName: !GetAtt ZSCCMacroFunction.Arn
      Principal: 'cloudformation.amazonaws.com'

  ZSCCMacro:
    Type: AWS::CloudFormation::Macro
    Properties:
      Name: 'ZSCC-Macro'
      Description: Zscaler CC Custom Macro
      FunctionName: !GetAtt ZSCCMacroFunction.Arn
Outputs:
  CloudConnectorMacroTemplateVersion:
    Description: Cloud Connector Macro Template Version
    Value: 1.0.0