Cartography EC2 route_tables ingest draft

route_tables.py

import logging
from typing import Dict
from typing import List

import boto3
import neo4j

from .util import get_botocore_config
from cartography.util import aws_handle_regions
from cartography.util import run_cleanup_job
from cartography.util import timeit

logger = logging.getLogger(__name__)


@timeit
@aws_handle_regions
def get_route_tables_data(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
    return client.describe_route_tables()['RouteTables']


@timeit
def load_route_tables(
    neo4j_session: neo4j.Session, data: List[Dict], region: str,
    aws_account_id: str, update_tag: int,
) -> None:
    ingest_route_tables = """
    UNWIND {route_tables} AS route_table

    MERGE (rt:Ec2RouteTable{id: route_table.RouteTableId})
    ON CREATE SET rt.firstseen = timestamp()
    SET rt.lastupdated = {update_tag},
    rt.route_table_id = route_table.RouteTableId
   
    MERGE (vpc:AWSVpc{id: route_table.VpcId})
    ON CREATE SET avpc.firstseen = timestamp()
    SET vpc.lastupdated = {update_tag}, vpc.vpcid = route_table.VpcId

    MERGE (account:AWSAccount{id: route_table.OwnerId})
    ON CREATE SET account.firstseen = timestamp()
    SET account.lastupdated = {update_tag}

    MERGE (account)-[r:RESOURCE]->(route_table)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}

    MERGE (vpc)-[r:RESOURCE]->(route_table)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}
   

    """

    neo4j_session.run(
        ingest_route_tables, route_tables=data, update_tag=update_tag,
        region=region, aws_account_id=aws_account_id,
    )


@timeit
def load_subnet_associations(
    neo4j_session: neo4j.Session, data: List[Dict], region: str,
    aws_account_id: str, update_tag: int,
) -> None:

    ingest_subnet_associations = """
    UNWIND {route_tables} AS route_table
    UNWIND route_table.Associations AS assoc

  
    MERGE (subnet:EC2Subnet{id: assoc.SubnetId})
    ON CREATE SET subnet.firstseen = timestamp()
    SET subnet.lastupdated = {update_tag}


    MERGE (subnet)-[r:ASSOCIATED]->(route_table)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}, r.main = assoc.Main,
    r.state = assoc.AssociationState.State,
    r.id = assoc.RouteTableAssociationId
    """

    neo4j_session.run(
        ingest_subnet_associations, route_tables=data, update_tag=update_tag,
        region=region, aws_account_id=aws_account_id,
    )


@timeit
def load_routes(
    neo4j_session: neo4j.Session, data: List[Dict], region: str,
    aws_account_id: str, update_tag: int,
) -> None:

    for route_table in data:
      # routes will be split to
      route_table['localRoutes'] = []
      route_table['InternetGatewayRoutes'] = []
      route_table['VPCPeeringRoutes'] = []
      route_table['NATGatewayRoutes'] = [] 
      route_table['TransitGatewayRoutes'] = [] 
      route_table['CarrierGatewayRoutes'] = []

      for route in route_table['Routes']:
        if 'DestinationCidrBlock' not in route:
          continue
        if route.get('GatewayId','') == 'local':
          route_table['localRoutes'].append(route)
        elif route.get('GatewayId','').find('igw-') == 0:
          route_table['InternetGatewayRoutes'].append(route)
        elif route.get('VpcPeeringConnectionId','').find('pcx-') == 0:
          route_table['VPCPeeringRoutes'].append(route)
        elif route.get('NatGatewayId','').find('nat-') == 0:
          route_table['NATGatewayRoutes'].append(route)
        elif route.get('TransitGatewayId','').find('tgw-') == 0:
          route_table['TransitGatewayRoutes'].append(route)        
        elif route.get('CarrierGatewayId','').find('cgw-') == 0:
          route_table['CarrierGatewayRoutes'].append(route)  

    ingest_local_routes = """
    UNWIND {route_tables} AS route_table
    UNWIND route_table.localRoutes AS route

    MERGE (ec2r:EC2Route:Route{id: route_table.RouteTableId + '|' + route.DestinationCidrBlock})
    ON CREATE SET ec2r.firstseen = timestamp()
    SET ec2r.lastupdated = {update_tag}, 
    ec2r.gateway_id  =  route.GatewayId,
    ec2r.destination_cidr_block = route.DestinationCidrBlock,
    ec2r.origin = route.Origin, ec2r.state = route.State

    MERGE (account:AWSAccount{id: route_table.OwnerId})
    ON CREATE SET account.firstseen = timestamp()
    SET account.lastupdated = {update_tag}  

    MERGE (account)-[r:RESOURCE]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}

    """
    neo4j_session.run(
        ingest_local_routes, route_tables=data, update_tag=update_tag,
        region=region, aws_account_id=aws_account_id,
    )

    ingest_igw_routes = """
    UNWIND {route_tables} AS route_table
    UNWIND route_table.InternetGatewayRoutes AS route

    MERGE (ec2r:EC2Route:Route{id: route_table.RouteTableId + '|' + route.DestinationCidrBlock})
    ON CREATE SET ec2r.firstseen = timestamp()
    SET ec2r.lastupdated = {update_tag}, 
    ec2r.gateway_id  =  route.GatewayId,
    ec2r.destination_cidr_block = route.DestinationCidrBlock,
    ec2r.origin = route.Origin, ec2r.state = route.State
  
    MERGE (igw:AWSInternetGateway{id: route.GatewayId})
    ON CREATE SET subnet.firstseen = timestamp()
    SET subnet.lastupdated = {update_tag}

    MERGE (igw)-[r:ASSOCIATED]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}, r.main = assoc.Main,
    r.state = assoc.AssociationState.State,
    r.id = assoc.RouteTableAssociationId    

    MERGE (account:AWSAccount{id: route_table.OwnerId})
    ON CREATE SET account.firstseen = timestamp()
    SET account.lastupdated = {update_tag}  

    MERGE (account)-[r:RESOURCE]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}

    """
   
    neo4j_session.run(
        ingest_igw_routes, route_tables=data, update_tag=update_tag,
        region=region, aws_account_id=aws_account_id,
    )

    ingest_pcx_routes = """
    UNWIND {route_tables} AS route_table
    UNWIND route_table.VPCPeeringRoutes AS route

    MERGE (ec2r:EC2Route:Route{id: route_table.RouteTableId + '|' + route.DestinationCidrBlock})
    ON CREATE SET ec2r.firstseen = timestamp()
    SET ec2r.lastupdated = {update_tag}, 
    ec2r.peeringconnection_id  =  route.VpcPeeringConnectionId,
    ec2r.destination_cidr_block = route.DestinationCidrBlock,
    ec2r.origin = route.Origin, ec2r.state = route.State
  
    MERGE (pcx:PeeringConnection{id: route.VpcPeeringConnectionId})
    ON CREATE SET subnet.firstseen = timestamp()
    SET subnet.lastupdated = {update_tag}

    MERGE (pcx)-[r:ASSOCIATED]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}, r.main = assoc.Main,
    r.state = assoc.AssociationState.State,
    r.id = assoc.RouteTableAssociationId    

    MERGE (account:AWSAccount{id: route_table.OwnerId})
    ON CREATE SET account.firstseen = timestamp()
    SET account.lastupdated = {update_tag}  

    MERGE (account)-[r:RESOURCE]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}


    """

    neo4j_session.run(
        ingest_pcx_routes, route_tables=data, update_tag=update_tag,
        region=region, aws_account_id=aws_account_id,
    )


    ingest_nat_routes = """
    UNWIND {route_tables} AS route_table
    UNWIND route_table.NATGatewayRoutes AS route

    MERGE (ec2r:EC2Route:Route{id: route_table.RouteTableId + '|' + route.DestinationCidrBlock})
    ON CREATE SET ec2r.firstseen = timestamp()
    SET ec2r.lastupdated = {update_tag}, 
    ec2r.natgateway_id  =  route.NatGatewayId,
    ec2r.destination_cidr_block = route.DestinationCidrBlock,
    ec2r.origin = route.Origin, ec2r.state = route.State

    MERGE (ngw:NatGateway{id: route.NatGatewayId})
    ON CREATE SET subnet.firstseen = timestamp()
    SET subnet.lastupdated = {update_tag}

    MERGE (ngw)-[r:ASSOCIATED]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}, r.main = assoc.Main,
    r.state = assoc.AssociationState.State,
    r.id = assoc.RouteTableAssociationId    

    MERGE (account:AWSAccount{id: route_table.OwnerId})
    ON CREATE SET account.firstseen = timestamp()
    SET account.lastupdated = {update_tag}  

    MERGE (account)-[r:RESOURCE]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}

    """


    neo4j_session.run(
        ingest_nat_routes, route_tables=data, update_tag=update_tag,
        region=region, aws_account_id=aws_account_id,
    )


    ingest_tgw_routes = """
    UNWIND {route_tables} AS route_table
    UNWIND route_table.TransitGatewayRoutes AS route

    MERGE (ec2r:EC2Route:Route{id: route_table.RouteTableId + '|' + route.DestinationCidrBlock})
    ON CREATE SET ec2r.firstseen = timestamp()
    SET ec2r.lastupdated = {update_tag}, 
    ec2r.transitgateway_id  =  route.TransitGatewayId,
    ec2r.destination_cidr_block = route.DestinationCidrBlock,
    ec2r.origin = route.Origin, ec2r.state = route.State

    MERGE (tgw:AWSTransitGateway{tgw_id: route.TransitGatewayId})
    ON CREATE SET subnet.firstseen = timestamp()
    SET subnet.lastupdated = {update_tag}

    MERGE (tgw)-[r:ASSOCIATED]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}, r.main = assoc.Main,
    r.state = assoc.AssociationState.State,
    r.id = assoc.RouteTableAssociationId    

    MERGE (account:AWSAccount{id: route_table.OwnerId})
    ON CREATE SET account.firstseen = timestamp()
    SET account.lastupdated = {update_tag}  

    MERGE (account)-[r:RESOURCE]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}

    """

    neo4j_session.run(
        ingest_tgw_routes, route_tables=data, update_tag=update_tag,
        region=region, aws_account_id=aws_account_id,
    )


    ingest_cgw_routes = """
    UNWIND {route_tables} AS route_table
    UNWIND route_table.CarrierGatewayRoutes AS route

    MERGE (ec2r:EC2Route:Route{id: route_table.RouteTableId + '|' + route.DestinationCidrBlock})
    ON CREATE SET ec2r.firstseen = timestamp()
    SET ec2r.lastupdated = {update_tag}, 
    ec2r.carriergateway_id  =  route.CarrierGatewayId,
    ec2r.destination_cidr_block = route.DestinationCidrBlock,
    ec2r.origin = route.Origin, ec2r.state = route.State
  
    MERGE (cgw:CarrierGatewayRoutes{id: route.CarrierGatewayId})
    ON CREATE SET subnet.firstseen = timestamp()
    SET subnet.lastupdated = {update_tag}

    MERGE (cgw)-[r:ASSOCIATED]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}, r.main = assoc.Main,
    r.state = assoc.AssociationState.State,
    r.id = assoc.RouteTableAssociationId    

    MERGE (account:AWSAccount{id: route_table.OwnerId})
    ON CREATE SET account.firstseen = timestamp()
    SET account.lastupdated = {update_tag}  

    MERGE (account)-[r:RESOURCE]->(route)
    ON CREATE SET r.firstseen = timestamp()
    SET r.lastupdated = {update_tag}



    """

    neo4j_session.run(
        ingest_cgw_routes, route_tables=data, update_tag=update_tag,
        region=region, aws_account_id=aws_account_id,
    )





@timeit
def cleanup_route_tables(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
    run_cleanup_job('aws_import_route_tables_cleanup.json', neo4j_session, common_job_parameters)


@timeit
def sync_route_tables(
    neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str],
    current_aws_account_id: str, update_tag: int, common_job_parameters: Dict,
) -> None:
    for region in regions:
        logger.debug("Syncing RouteTables for region '%s' in account '%s'.", region, current_aws_account_id)
        data = get_route_tables_data(boto3_session, region)
        load_route_tables(neo4j_session, data, region, current_aws_account_id, update_tag)
        load_subnet_associations(neo4j_session, data, region, current_aws_account_id, update_tag)
        load_routes(neo4j_session, data, region, current_aws_account_id, update_tag)
    cleanup_route_tables(neo4j_session, common_job_parameters)