import gzip
import json
import base64
import boto3
import botocore.exceptions
import ipaddress
import os
import math

SPOKE_ACCOUNT_IDS = os.environ["SPOKE_ACCOUNT_IDS"] 
CURRENT_ACCOUNT = os.environ["CURRENT_ACCOUNT"]
DDB_NAME = os.environ["DDB_NAME"]

def scanddb ():
    ddb = boto3.client('dynamodb')
    SubnetAZIdList = []
    table={}
    table = ddb.scan(TableName = DDB_NAME)
    for i in table['Items']:
        SubnetAZIdDict={}
        SubnetAZIdDict['AvailabilityZoneId']=i['AvailabilityZoneId']['S']
        SubnetAZIdDict['CidrBlock']=i['CidrBlock']['S']
        SubnetAZIdDict['SubnetId']=i['SubnetId']['S']
        SubnetAZIdList.append(SubnetAZIdDict)
    return SubnetAZIdList
    
def getSubnetIdFromENI(eni, account_id=None):
    if not account_id:
      try: 
          ec2 = boto3.client('ec2')
          response = ec2.describe_network_interfaces(
              NetworkInterfaceIds=[eni]  
          )
          # print(response)
          subnetId = response['NetworkInterfaces'][0]['SubnetId']
          # print(subnetId)
          return subnetId
      except botocore.exceptions.ClientError as e:
          print(e)
          pass
    else: 
      try:
          sts_connection = boto3.client('sts')
          arn_tuple = ("arn:aws:iam::",account_id,":role/DtazExecRole")
          acct_b = sts_connection.assume_role(
              RoleArn="".join(arn_tuple),
              RoleSessionName="cross_acct_lambda"
          )
          ACCESS_KEY = acct_b['Credentials']['AccessKeyId']
          SECRET_KEY = acct_b['Credentials']['SecretAccessKey']
          SESSION_TOKEN = acct_b['Credentials']['SessionToken']
          
          ec2 = boto3.client(
                'ec2',
                aws_access_key_id=ACCESS_KEY,
                aws_secret_access_key=SECRET_KEY,
                aws_session_token=SESSION_TOKEN
          )
          response = ec2.describe_network_interfaces(
                NetworkInterfaceIds=[eni]  
            )
          subnetId = response['NetworkInterfaces'][0]['SubnetId']
          # print(subnetId)
          return subnetId
      except botocore.exceptions.ClientError as e:
          print(e)
          pass

          
def addressInNetwork(ip_address,subnet_cidr):
    try:
        IPAddress = ipaddress.ip_address(ip_address)
        SubnetCidr = ipaddress.ip_network(subnet_cidr)
    except ValueError:
        IPAddress = None 
        SubnetCidr = None

    if (IPAddress or SubnetCidr) == None:
        #print('skip for bad IP/CIDR')
        pass
    else:
        AddressInNetwork = IPAddress in SubnetCidr
        return AddressInNetwork
        
def lambda_handler(event, context):
    Response = {}
    Response['event']={}
    cw_data = ""
    srDestBytelist=[]
    
    TenPnine = math.pow(10,9)
    if SPOKE_ACCOUNT_IDS == "NoValue":
      cw_data = event['awslogs']['data']
      compressed_payload = base64.b64decode(cw_data)
      uncompressed_payload = gzip.decompress(compressed_payload)
      payload = json.loads(uncompressed_payload)
      ENI_id = payload['logEvents'][0]['extractedFields']['interface_id']
      # print(ENI_id)
      subnet_id = getSubnetIdFromENI(ENI_id)
      # print("Subnet ID:", subnet_id)
      log_event = payload['logEvents']

      
    if SPOKE_ACCOUNT_IDS != "NoValue":
      cw_data = event['Records'][0]['kinesis']['data']
      compressed_payload = base64.b64decode(cw_data)
      uncompressed_payload = gzip.decompress(compressed_payload)
      payload = json.loads(uncompressed_payload)
      ENI_id = payload['logEvents'][0]['extractedFields']['interface_id']
      owner_id = payload['owner']
      # print(ENI_id)
      subnet_id = getSubnetIdFromENI(ENI_id,owner_id)
      # print("Subnet ID:", subnet_id)
      log_event = payload['logEvents']

    for stream in log_event:
      srDestByteDict={}
      srDestByteDict['srcaddr'] = stream['extractedFields']['srcaddr']
      srDestByteDict['dstaddr'] = stream['extractedFields']['dstaddr']
      srDestByteDict['bytes'] = stream['extractedFields']['bytes']
      srDestByteDict['ENIId'] = ENI_id
      srDestByteDict['subnetid'] = subnet_id
      srDestBytelist.append(srDestByteDict)
      # print(json.dumps(srDestBytelist))
      
    subnet_list_from_ddb = scanddb()

    
    for item in srDestBytelist: 
       for az in subnet_list_from_ddb:
          # check if the source address is in one CIDR Block 
          if addressInNetwork(item['srcaddr'],az['CidrBlock']):
              # check the validation of the ENIs by subnet ID to avoid CIDR block overlapping 
              if item['subnetid'] == az['SubnetId']:
                  azSrcID = az['AvailabilityZoneId']
                  for az in subnet_list_from_ddb:
                      # Check the Destination IP address is in one CIDR Block 
                      if addressInNetwork(item['dstaddr'],az['CidrBlock']):
                          azDestID = az['AvailabilityZoneId']
                          # Check if the Destination IP address adn source IP address are from same AZ. 
                          if azSrcID != azDestID:
                              Response['event']['srIp'] = item['srcaddr']
                              Response['event']['destIp'] = item['dstaddr']
                              Response['event']['bytes'] = int(item['bytes'])
                              Response['event']['ENI'] = item['ENIId']
                              temp_data = Response['event']['bytes'] 

                              if temp_data > TenPnine:
                                up = math.ceil(temp_data/TenPnine)
                                for i in range(up):
                                    temp_data = temp_data - TenPnine 
                                    if temp_data >= 0:
                                        Response['event']['bytes'] = TenPnine - 1
                                    else:
                                        Response['event']['bytes'] = TenPnine + temp_data + up
                                    print(json.dumps(Response))
                              else:
                                print(json.dumps(Response))