Skip to content

Instantly share code, notes, and snippets.

@krisek

krisek/route_tables.py

Created March 21, 2021 21:49
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save krisek/81e6fd1d654214d10b1b8c565ee61d23 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save krisek/81e6fd1d654214d10b1b8c565ee61d23 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. krisek created this gist Mar 21, 2021.
    353 changes: 353 additions & 0 deletions route_tables.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,353 @@
    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)