cartography 0.102.0rc1__py3-none-any.whl → 0.103.0__py3-none-any.whl

cartography/intel/aws/ec2/network_interfaces.py

@@ -8,20 +8,28 @@ from typing import List
 import boto3
 import neo4j
 
-from .util import get_botocore_config
 from cartography.client.core.tx import load
 from cartography.graph.job import GraphJob
 from cartography.models.aws.ec2.networkinterfaces import EC2NetworkInterfaceSchema
-from cartography.models.aws.ec2.privateip_networkinterface import EC2PrivateIpNetworkInterfaceSchema
-from cartography.models.aws.ec2.securitygroup_networkinterface import EC2SecurityGroupNetworkInterfaceSchema
-from cartography.models.aws.ec2.subnet_networkinterface import EC2SubnetNetworkInterfaceSchema
+from cartography.models.aws.ec2.privateip_networkinterface import (
+    EC2PrivateIpNetworkInterfaceSchema,
+)
+from cartography.models.aws.ec2.securitygroup_networkinterface import (
+    EC2SecurityGroupNetworkInterfaceSchema,
+)
+from cartography.models.aws.ec2.subnet_networkinterface import (
+    EC2SubnetNetworkInterfaceSchema,
+)
 from cartography.util import aws_handle_regions
 from cartography.util import timeit
 
+from .util import get_botocore_config
+
 logger = logging.getLogger(__name__)
 
 Ec2NetworkData = namedtuple(
-    "Ec2NetworkData", [
+    "Ec2NetworkData",
+    [
         "network_interface_list",
         "private_ip_list",
         "sg_list",
@@ -32,16 +40,26 @@ Ec2NetworkData = namedtuple(
 
 @timeit
 @aws_handle_regions
-def get_network_interface_data(boto3_session: boto3.session.Session, region: str) -> List[Dict[str, Any]]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
-    paginator = client.get_paginator('describe_network_interfaces')
+def get_network_interface_data(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict[str, Any]]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    paginator = client.get_paginator("describe_network_interfaces")
     subnets: List[Dict] = []
     for page in paginator.paginate():
-        subnets.extend(page['NetworkInterfaces'])
+        subnets.extend(page["NetworkInterfaces"])
     return subnets
 
 
-def transform_network_interface_data(data_list: List[Dict[str, Any]], region: str) -> Ec2NetworkData:
+def transform_network_interface_data(
+    data_list: List[Dict[str, Any]],
+    region: str,
+) -> Ec2NetworkData:
     network_interface_list = []
     private_ip_list = []
     sg_list = []
@@ -52,45 +70,52 @@ def transform_network_interface_data(data_list: List[Dict[str, Any]], region: st
         # https://aws.amazon.com/premiumsupport/knowledge-center/elb-find-load-balancer-IP/
         elb_v1_id = None
         elb_v2_id = None
-        elb_match = re.match(r'^ELB (?:net|app)/([^\/]+)\/(.*)', network_interface.get('Description', ''))
+        elb_match = re.match(
+            r"^ELB (?:net|app)/([^\/]+)\/(.*)",
+            network_interface.get("Description", ""),
+        )
         if elb_match:
-            elb_v1_id = f'{elb_match[1]}-{elb_match[2]}.elb.{region}.amazonaws.com'
+            elb_v1_id = f"{elb_match[1]}-{elb_match[2]}.elb.{region}.amazonaws.com"
         else:
-            elb_match = re.match(r'^ELB (.*)', network_interface.get('Description', ''))
+            elb_match = re.match(r"^ELB (.*)", network_interface.get("Description", ""))
             if elb_match:
                 elb_v2_id = elb_match[1]
         # TODO issue #1024 change this to arn when ready
-        network_interface_id = network_interface['NetworkInterfaceId']
+        network_interface_id = network_interface["NetworkInterfaceId"]
         network_interface_list.append(
             {
-                'Id': network_interface_id,
-                'NetworkInterfaceId': network_interface['NetworkInterfaceId'],
-                'Description': network_interface['Description'],
-                'InstanceId': network_interface.get('Attachment', {}).get('InstanceId'),
-                'InterfaceType': network_interface['InterfaceType'],
-                'MacAddress': network_interface['MacAddress'],
-                'PrivateDnsName': network_interface.get('PrivateDnsName'),
-                'PrivateIpAddress': network_interface['PrivateIpAddress'],
-                'PublicIp': network_interface.get('Association', {}).get('PublicIp'),
-                'RequesterId': network_interface.get('RequesterId'),
-                'RequesterManaged': network_interface['RequesterManaged'],
-                'SourceDestCheck': network_interface['SourceDestCheck'],
-                'Status': network_interface['Status'],
-                'SubnetId': network_interface['SubnetId'],
-                'ElbV1Id': elb_v1_id,
-                'ElbV2Id': elb_v2_id,
+                "Id": network_interface_id,
+                "NetworkInterfaceId": network_interface["NetworkInterfaceId"],
+                "Description": network_interface["Description"],
+                "InstanceId": network_interface.get("Attachment", {}).get("InstanceId"),
+                "InterfaceType": network_interface["InterfaceType"],
+                "MacAddress": network_interface["MacAddress"],
+                "PrivateDnsName": network_interface.get("PrivateDnsName"),
+                "PrivateIpAddress": network_interface["PrivateIpAddress"],
+                "PublicIp": network_interface.get("Association", {}).get("PublicIp"),
+                "RequesterId": network_interface.get("RequesterId"),
+                "RequesterManaged": network_interface["RequesterManaged"],
+                "SourceDestCheck": network_interface["SourceDestCheck"],
+                "Status": network_interface["Status"],
+                "SubnetId": network_interface["SubnetId"],
+                "ElbV1Id": elb_v1_id,
+                "ElbV2Id": elb_v2_id,
             },
         )
-        if network_interface.get('PrivateIpAddresses'):
-            for private_ip_address in network_interface['PrivateIpAddresses']:
+        if network_interface.get("PrivateIpAddresses"):
+            for private_ip_address in network_interface["PrivateIpAddresses"]:
                 private_ip_list.append(
                     {
-                        'Id': f"{network_interface['NetworkInterfaceId']}:{private_ip_address['PrivateIpAddress']}",
-                        'NetworkInterfaceId': network_interface['NetworkInterfaceId'],
-                        'IpOwnerId': private_ip_address.get('Association', {}).get('IpOwnerId'),
-                        'Primary': private_ip_address['Primary'],
-                        'PrivateIpAddress': private_ip_address['PrivateIpAddress'],
-                        'PublicIp': private_ip_address.get('Association', {}).get('PublicIp'),
+                        "Id": f"{network_interface['NetworkInterfaceId']}:{private_ip_address['PrivateIpAddress']}",
+                        "NetworkInterfaceId": network_interface["NetworkInterfaceId"],
+                        "IpOwnerId": private_ip_address.get("Association", {}).get(
+                            "IpOwnerId",
+                        ),
+                        "Primary": private_ip_address["Primary"],
+                        "PrivateIpAddress": private_ip_address["PrivateIpAddress"],
+                        "PublicIp": private_ip_address.get("Association", {}).get(
+                            "PublicIp",
+                        ),
                     },
                 )
 
@@ -98,19 +123,19 @@ def transform_network_interface_data(data_list: List[Dict[str, Any]], region: st
             for group in network_interface["Groups"]:
                 sg_list.append(
                     {
-                        'GroupId': group['GroupId'],
-                        'NetworkInterfaceId': network_interface_id,
+                        "GroupId": group["GroupId"],
+                        "NetworkInterfaceId": network_interface_id,
                     },
                 )
 
-        subnet_id = network_interface.get('SubnetId')
+        subnet_id = network_interface.get("SubnetId")
         if subnet_id:
             subnet_list.append(
                 {
-                    'NetworkInterfaceId': network_interface_id,
-                    'SubnetId': subnet_id,
-                    'ElbV1Id': elb_v1_id,
-                    'ElbV2Id': elb_v2_id,
+                    "NetworkInterfaceId": network_interface_id,
+                    "SubnetId": subnet_id,
+                    "ElbV1Id": elb_v1_id,
+                    "ElbV2Id": elb_v2_id,
                 },
             )
 
@@ -124,11 +149,11 @@ def transform_network_interface_data(data_list: List[Dict[str, Any]], region: st
 
 @timeit
 def load_network_interfaces(
-        neo4j_session: neo4j.Session,
-        data: List[Dict[str, Any]],
-        region: str,
-        aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     logger.info(f"Loading {len(data)} network interfaces in {region}.")
     load(
@@ -143,11 +168,11 @@ def load_network_interfaces(
 
 @timeit
 def load_private_ip_network_interface(
-        neo4j_session: neo4j.Session,
-        data: List[Dict[str, Any]],
-        region: str,
-        aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     """
     Private IPs as known by describe-network-interfaces.
@@ -165,11 +190,11 @@ def load_private_ip_network_interface(
 
 @timeit
 def load_security_group_network_interface(
-        neo4j_session: neo4j.Session,
-        data: List[Dict[str, Any]],
-        region: str,
-        aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     """
     Security groups as known by describe-network-interfaces.
@@ -187,11 +212,11 @@ def load_security_group_network_interface(
 
 @timeit
 def load_subnet_network_interface(
-        neo4j_session: neo4j.Session,
-        data: List[Dict[str, Any]],
-        region: str,
-        aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     """
     Subnets as known by describe-network-interfaces.
@@ -208,38 +233,72 @@ def load_subnet_network_interface(
 
 
 def load_network_data(
-        neo4j_session: neo4j.Session,
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
-        network_interface_list: List[Dict[str, Any]],
-        private_ip_list: List[Dict[str, Any]],
-        subnet_list: List[Dict[str, Any]],
-        sg_list: List[Dict[str, Any]],
+    neo4j_session: neo4j.Session,
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
+    network_interface_list: List[Dict[str, Any]],
+    private_ip_list: List[Dict[str, Any]],
+    subnet_list: List[Dict[str, Any]],
+    sg_list: List[Dict[str, Any]],
 ) -> None:
-    load_network_interfaces(neo4j_session, network_interface_list, region, current_aws_account_id, update_tag)
-    load_private_ip_network_interface(neo4j_session, private_ip_list, region, current_aws_account_id, update_tag)
-    load_subnet_network_interface(neo4j_session, subnet_list, region, current_aws_account_id, update_tag)
-    load_security_group_network_interface(neo4j_session, sg_list, region, current_aws_account_id, update_tag)
+    load_network_interfaces(
+        neo4j_session,
+        network_interface_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_private_ip_network_interface(
+        neo4j_session,
+        private_ip_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_subnet_network_interface(
+        neo4j_session,
+        subnet_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_security_group_network_interface(
+        neo4j_session,
+        sg_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
 
 
 @timeit
-def cleanup_network_interfaces(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    GraphJob.from_node_schema(EC2NetworkInterfaceSchema(), common_job_parameters).run(neo4j_session)
-    GraphJob.from_node_schema(EC2PrivateIpNetworkInterfaceSchema(), common_job_parameters).run(neo4j_session)
+def cleanup_network_interfaces(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
+    GraphJob.from_node_schema(EC2NetworkInterfaceSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+    GraphJob.from_node_schema(
+        EC2PrivateIpNetworkInterfaceSchema(),
+        common_job_parameters,
+    ).run(neo4j_session)
 
 
 @timeit
 def sync_network_interfaces(
-        neo4j_session: neo4j.Session,
-        boto3_session: boto3.session.Session,
-        regions: List[str],
-        current_aws_account_id: str,
-        update_tag: int,
-        common_job_parameters: Dict,
+    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.info(f"Syncing EC2 network interfaces for region '{region}' in account '{current_aws_account_id}'.")
+        logger.info(
+            f"Syncing EC2 network interfaces for region '{region}' in account '{current_aws_account_id}'.",
+        )
         data = get_network_interface_data(boto3_session, region)
         ec2_network_data = transform_network_interface_data(data, region)
         load_network_data(

cartography/intel/aws/ec2/reserved_instances.py

@@ -6,30 +6,43 @@ import boto3
 import neo4j
 from botocore.exceptions import ClientError
 
-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
 
+from .util import get_botocore_config
+
 logger = logging.getLogger(__name__)
 
 
 @timeit
 @aws_handle_regions
-def get_reserved_instances(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
+def get_reserved_instances(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
     try:
-        reserved_instances = client.describe_reserved_instances()['ReservedInstances']
+        reserved_instances = client.describe_reserved_instances()["ReservedInstances"]
     except ClientError as e:
-        logger.warning(f"Failed retrieve reserved instances for region - {region}. Error - {e}")
+        logger.warning(
+            f"Failed retrieve reserved instances for region - {region}. Error - {e}",
+        )
         raise
     return reserved_instances
 
 
 @timeit
 def load_reserved_instances(
-        neo4j_session: neo4j.Session, data: List[Dict], region: str,
-        current_aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     ingest_reserved_instances = """
     UNWIND $reserved_instances_list as res
@@ -48,8 +61,8 @@ def load_reserved_instances(
     """
 
     for r_instance in data:
-        r_instance['Start'] = str(r_instance['Start'])
-        r_instance['End'] = str(r_instance['End'])
+        r_instance["Start"] = str(r_instance["Start"])
+        r_instance["End"] = str(r_instance["End"])
 
     neo4j_session.run(
         ingest_reserved_instances,
@@ -61,9 +74,12 @@ def load_reserved_instances(
 
 
 @timeit
-def cleanup_reserved_instances(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
+def cleanup_reserved_instances(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
     run_cleanup_job(
-        'aws_import_reserved_instances_cleanup.json',
+        "aws_import_reserved_instances_cleanup.json",
         neo4j_session,
         common_job_parameters,
     )
@@ -71,12 +87,25 @@ def cleanup_reserved_instances(neo4j_session: neo4j.Session, common_job_paramete
 
 @timeit
 def sync_ec2_reserved_instances(
-        neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str],
-        current_aws_account_id: str,
-        update_tag: int, common_job_parameters: Dict,
+    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 reserved instances for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.debug(
+            "Syncing reserved instances for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_reserved_instances(boto3_session, region)
-        load_reserved_instances(neo4j_session, data, region, current_aws_account_id, update_tag)
+        load_reserved_instances(
+            neo4j_session,
+            data,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_reserved_instances(neo4j_session, common_job_parameters)