import json
import logging
import boto3
import cfnresponse

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

def get_vpn_attachment_ids(vpn1_id, vpn2_id, stackName):
    try:
      myDict = {}
      client = boto3.client('ec2')
      vpn = client.describe_vpn_connections(VpnConnectionIds=[vpn1_id,vpn2_id])['VpnConnections']
      tgw = client.describe_transit_gateway_attachments()['TransitGatewayAttachments']
      rtb = client.describe_transit_gateway_route_tables()["TransitGatewayRouteTables"]
      logger.info(vpn)
      for index in range(len(vpn)):
        mylist = []
        for vgwTelemetry in vpn[index]['VgwTelemetry']:
          mylist.append(vgwTelemetry['OutsideIpAddress'])
        myDict['vpn'+str(index)+'OutsideIps']=mylist
      for dictionary in tgw:
        if dictionary["Tags"]:
          for tagsdictionary in dictionary["Tags"]:
            if stackName+"-HRVPC" in tagsdictionary['Value']:
              myDict["HRVPC_tgw_attachment_id"] = dictionary['TransitGatewayAttachmentId']
            if stackName+"-ENGVPC" in tagsdictionary['Value']:
              myDict["ENGVPC_tgw_attachment_id"] = dictionary['TransitGatewayAttachmentId']
        elif dictionary["ResourceId"] == vpn1_id:
          myDict["vpn1_tgw_attachment_id"] = dictionary['TransitGatewayAttachmentId']
        elif dictionary["ResourceId"] == vpn2_id:
          myDict["vpn2_tgw_attachment_id"] = dictionary['TransitGatewayAttachmentId']
      for dictionary in rtb:
        if dictionary["Tags"]:
          for tagsdictionary in dictionary["Tags"]:
            if stackName+"-edge" in tagsdictionary['Value']:
              myDict["edge_tgw_rtb_id"] = dictionary['TransitGatewayRouteTableId']
            if stackName+"-HRVPC" in tagsdictionary['Value']:
              myDict["HRVPC_tgw_rtb_id"] = dictionary['TransitGatewayRouteTableId']
            if stackName+"-ENGVPC" in tagsdictionary['Value']:
              myDict["ENGVPC_tgw_rtb_id"] = dictionary['TransitGatewayRouteTableId']
    except Exception as e:
      logger.info('get vpn tgw attachment id failure: {}'.format(e))

    return myDict

def lambda_handler(event, context):
    logger.info('got event {}'.format(event))
    responseData = {}
    
    if event['RequestType'] == 'Create':
      responseData = get_vpn_attachment_ids(event['ResourceProperties']['vpn1_id'],event['ResourceProperties']['vpn2_id'],event['ResourceProperties']['stackName'])
    
    else: # delete / update
      rs = event['PhysicalResourceId'] 
      responseData['TransitGatewayAttributes'] = rs
    
    logger.info('responseData {}'.format(responseData))
    cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData)