from diagrams import Cluster, Diagram, Edge from diagrams.aws.compute import EC2 from diagrams.aws.network import Endpoint from diagrams.onprem.compute import Server from diagrams.onprem.network import Internet from botocore.exceptions import ClientError import math from pingcap.pe.flash.util import get_boto_session, get_name_by_datetime # Basic tools def get_vpc_default_rt(vpc_id: str, profile=None) -> dict: """ Get main RT in VPC """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_route_tables( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, { "Name": "association.main", "Values": [ "true", ], }, ] ) result = response["RouteTables"][0] return result def get_short_name(name) -> str: """ " Cloudformation always generate extreme long object name concated by '-' """ name_long = name.split("-") name_short = "-".join(name_long[: math.ceil(len(name_long) / 1.2)]) return name_short def get_12_digits_account_number(profile=None) -> str: session = get_boto_session(profile) sts = session.client("sts") return sts.get_caller_identity()["Arn"].split(":")[4] # Security reports format def merge_sec_line(record: dict, sec_line: dict) -> None: """ Merge security record into sec_line and the sec_line. 5 fields in sec_line: key: "resource_type,resource_tag,resource_name" value: "related_resource,best_practice,current_practice, advice" """ ret = {} for k in [*record]: if k in sec_line: sec_line[k] = record[k] else: ret = sec_line.update(record) sec_line = ret # For IAM probe def get_groups_by_user(username: str, profile=None) -> list: session = get_boto_session(profile) iam = session.client("iam") response = iam.list_groups_for_user(UserName=username, MaxItems=100) result = response["Groups"] while True: if "Marker" in response: response = iam.list_groups_for_user( UserName=username, MaxItems=100, Marker=response["Marker"] ) result = result + response["Groups"] else: break return result def get_all_iam_users(profile=None) -> list: session = get_boto_session(profile) iam = session.client("iam") response = iam.list_users(MaxItems=100) result = response["Users"] while True: if "Marker" in response: response = iam.list_users(MaxItems=100, Marker=response["Marker"]) result = result + response["Users"] else: break return result def get_inline_user_policies(username: str, profile=None) -> list: session = get_boto_session(profile) iam = session.client("iam") response = iam.list_user_policies(UserName=username, MaxItems=100) result = response["PolicyNames"] while True: if "Marker" in response: response = iam.list_user_policies( UserName=username, MaxItems=100, Marker=response["Marker"] ) result = result + response["PolicyNames"] else: break return result def get_inline_group_policies(groupname: str, profile=None) -> list: session = get_boto_session(profile) iam = session.client("iam") response = iam.list_group_policies(GroupName=groupname, MaxItems=100) result = response["PolicyNames"] while True: if "Marker" in response: response = iam.list_group_policies( GroupName=groupname, MaxItems=100, Marker=response["Marker"] ) result = result + response["PolicyNames"] else: break return result def get_inline_role_policies(rolename: str, profile=None) -> list: session = get_boto_session(profile) iam = session.client("iam") response = iam.list_role_policies(RoleName=rolename, MaxItems=100) result = response["PolicyNames"] while True: if "Marker" in response: response = iam.list_role_policies( RoleName=rolename, MaxItems=100, Marker=response["Marker"] ) result = result + response["PolicyNames"] else: break return result def get_all_iam_groups(profile=None): session = get_boto_session(profile) iam = session.client("iam") response = iam.list_groups(MaxItems=100) result = response["Groups"] while True: if "Marker" in response: response = iam.list_groups(MaxItems=100, Marker=response["Marker"]) result = result + response["Groups"] else: break return result def get_all_iam_roles(profile=None): session = get_boto_session(profile) iam = session.client("iam") response = iam.list_roles(MaxItems=100) result = response["Roles"] while True: if "Marker" in response: response = iam.list_roles(MaxItems=100, Marker=response["Marker"]) result = result + response["Roles"] else: break return result def get_local_iam_policies(profile=None): session = get_boto_session(profile) iam = session.client("iam") response = iam.list_policies(Scope="Local", OnlyAttached=True, MaxItems=100) result = response["Policies"] while True: if "Marker" in response: response = iam.list_policies( Scope="Local", OnlyAttached=True, MaxItems=100, Marker=response["Marker"], ) result = result + response["Policies"] else: break return result def get_aws_iam_policies(profile=None): session = get_boto_session(profile) iam = session.client("iam") response = iam.list_policies(Scope="AWS", OnlyAttached=True, MaxItems=100) result = response["Policies"] while True: if "Marker" in response: response = iam.list_policies( Scope="Local", OnlyAttached=True, MaxItems=100, Marker=response["Marker"], ) result = result + response["Policies"] else: break return result def get_attached_user_policies(username: str, profile=None) -> list: session = get_boto_session(profile) iam = session.client("iam") response = iam.list_attached_user_policies(UserName=username, MaxItems=100) result = response["AttachedPolicies"] while True: if "Marker" in response: response = iam.list_attached_user_policies( UserName=username, MaxItems=100, Marker=response["Marker"] ) result = result + response["AttachedPolicies"] else: break return result def get_attached_group_policies(groupname: str, profile=None) -> list: session = get_boto_session(profile) iam = session.client("iam") response = iam.list_attached_group_policies(GroupName=groupname, MaxItems=100) result = response["AttachedPolicies"] while True: if "Marker" in response: response = iam.list_attached_group_policies( GroupName=groupname, MaxItems=100, Marker=response["Marker"] ) result = result + response["AttachedPolicies"] else: break return result def get_attached_role_policies(rolename: str, profile=None) -> list: session = get_boto_session(profile) iam = session.client("iam") response = iam.list_attached_role_policies(RoleName=rolename, MaxItems=100) result = response["AttachedPolicies"] while True: if "Marker" in response: response = iam.list_attached_role_policies( RoleName=rolename, MaxItems=100, Marker=response["Marker"] ) result = result + response["AttachedPolicies"] else: break return result # For Region probe def get_non_az_level_resource_usages(profile=None) -> dict: """ Non-AZ scoped resources are listed here in one shot. Subnets are also AZ related. """ serv_list = { "s3": describe_buckets(profile), "sqs": describe_region_sqs(profile), "efs": describe_region_efs_fs(profile), "autoscaling": describe_region_asgs(profile), "snapshot": describe_region_snapshots(profile), "ami": describe_region_amis(profile), } return serv_list def get_nametag_value(source, name_hint=None, default=None) -> str: """ 1. Name tag 2. Native object name 3. Provided default name """ if "Tags" in source: for o in source["Tags"]: if o["Key"] == "Name": return o["Value"] if name_hint and name_hint in source and source[name_hint] != None: return source[name_hint] else: return default def map_ebs_id_to_asg_name(all_ebs_volumes, all_asg_instances) -> dict: """ returns {ebs_id: asg_name} """ ebs_to_asg = {} for vol in all_ebs_volumes: if "Attachments" in vol: for attach in vol["Attachments"]: for instance in all_asg_instances: if attach["InstanceId"] == instance["InstanceId"]: ebs_to_asg[attach["VolumeId"]] = instance[ "AutoScalingGroupName" ] return ebs_to_asg def map_instance_id_to_sgs_list(instance_to_nics, nic_to_sgs) -> dict: """ {:(set)} """ instance_to_sgs = {} sgs = set() for i in [*instance_to_nics]: sgs = set() for nic in instance_to_nics[i]: for sg in nic_to_sgs[nic]: sgs.add(sg) instance_to_sgs[i] = sgs return instance_to_sgs def map_rt_id_to_route_list(route_tables: list) -> dict: """ return: { :[] } """ rt_to_routes = {} for rt in route_tables: route_count = 0 routes = [] for route in rt["Routes"]: route_count += 1 keys = [*route] destination = "" target = "" for k in keys: if k.startswith("Destination"): destination += ( "To " + k.replace("Destination", "") + ": " + route[k] + ", " ) elif k.endswith("Id"): target += ( "via " + k.replace("Destination", "") + ": " + route[k] + ", " ) if route["Origin"] == "EnableVgwRoutePropagation": propagated = "yes" else: propagated = "no" state = route["State"] routes += [ "#" + str(route_count) + ": " + destination + target + "propagated: " + propagated + ", " + state ] rt_to_routes[rt["RouteTableId"]] = routes return rt_to_routes def map_sg_id_to_permissions(security_groups) -> dict: """ returns: { :{ 'in': [] 'out': [] } } """ sg_to_permissions = {} for sg in security_groups: sg_to_permissions[sg["GroupId"]] = {"in": [], "out": []} in_permission_idx = 0 for p in sg["IpPermissions"]: # print(p) in_permission_idx += 1 cidr = "" ip_protocol = p["IpProtocol"] if ip_protocol == "-1": ipp_info = "All" elif "FromPort" in p: ipp_info = ip_protocol + "-ports " if p["FromPort"] == p["ToPort"]: if p["FromPort"] == -1: ipp_info += "All" else: ipp_info += "(" + str(p["FromPort"]) + ")" else: ipp_info += "(" + str(p["FromPort"]) + "-" + str(p["ToPort"]) + ")" else: ipp_info = ip_protocol sgr = "" for ugp in p["UserIdGroupPairs"]: sgr += ( extract_sg_name_from_sgs_list(ugp["GroupId"], security_groups) + ", " ) keys = [*p] for k in keys: if k.endswith("Ranges"): for range in p[k]: cidr += range["Cidr" + k.replace("Ranges", "")] + ", " inbound_permission = "In: " + cidr + sgr + "Protocol: " + ipp_info sg_to_permissions[sg["GroupId"]]["in"] += [inbound_permission] out_permission_idx = 0 for p in sg["IpPermissionsEgress"]: out_permission_idx += 1 cidr = "" ip_protocol = p["IpProtocol"] if ip_protocol == "-1": ipp_info = "All" else: ipp_info = ip_protocol + "-ports " if p["FromPort"] == p["ToPort"]: if p["FromPort"] == -1: ipp_info += "All" else: ipp_info += "(" + str(p["FromPort"]) + ")" else: ipp_info += "(" + str(p["FromPort"]) + "-" + str(p["ToPort"]) + ")" sgr = "" for ugp in p["UserIdGroupPairs"]: sgr += ( extract_sg_name_from_sgs_list(ugp["GroupId"], security_groups) + ", " ) keys = [*p] for k in keys: if k.endswith("Ranges"): for range in p[k]: cidr += range["Cidr" + k.replace("Ranges", "")] + ", " outbound_permission = "Out: " + cidr + sgr + "Protocol: " + ipp_info sg_to_permissions[sg["GroupId"]]["out"] += [outbound_permission] return sg_to_permissions def map_az_name_to_subnet_list(subnets) -> dict: """ {:[]} """ az_to_subnets = {subnet["AvailabilityZone"]: [] for subnet in subnets} for subnet in subnets: k = subnet["AvailabilityZone"] v = subnet ov = az_to_subnets[k] az_to_subnets[k] = ov + [v] return az_to_subnets # Simple describers def describe_buckets(profile: str) -> list: session = get_boto_session(profile) s3 = session.client("s3") res = s3.list_buckets() if "Buckets" in res: return res["Buckets"] else: return [] def describe_region_sqs(profile) -> list: session = get_boto_session(profile) sqs = session.client("sqs") response = sqs.list_queues(MaxResults=100) result = None if "QueueUrls" in response: result = response["QueueUrls"] while True: if "NextToken" in response: response = sqs.list_queues(MaxResults=100, NextToken=response["NextToken"]) result = result + response["QueueUrls"] else: break if result and len(result) > 0: return result else: return [] def describe_region_vpcs(profile) -> list: session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_vpcs(MaxResults=100) result = response["Vpcs"] while True: if "NextToken" in response: response = ec2.describe_vpcs( MaxResults=100, NextToken=response["NextToken"] ) result = result + response["Vpcs"] else: break return result def describe_region_amis(profile) -> list: """ Return all AMIs in region. """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_images(Owners=["self"]) result = response["Images"] return result def describe_region_snapshots(profile) -> list: """ Return all snapshots in region. """ account_number = get_12_digits_account_number(profile) session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_snapshots( Filters=[ { "Name": "status", "Values": [ "completed", ], }, ], OwnerIds=[account_number], MaxResults=100, ) result = response["Snapshots"] while True: if "NextToken" in response: response = response = ec2.describe_snapshots( Filters=[ { "Name": "status", "Values": [ "completed", ], }, ], OwnerIds=[account_number], MaxResults=100, NextToken=response["NextToken"], ) result = result + response["Snapshots"] else: break return result def describe_region_efs_fs(profile) -> list: """ Regional EFS file systems """ session = get_boto_session(profile) efs = session.client("efs") response = efs.describe_file_systems(MaxItems=100) result = response["FileSystems"] while True: if "Marker" in response: response = efs.describe_file_systems(MaxItems=100, Marker=response["Maker"]) result = result + response["FileSystems"] else: break if result and len(result) > 0: return result else: return [] def describe_region_asgs(profile) -> list: """ Return all ASG groups in region. """ session = get_boto_session(profile) asg = session.client("autoscaling") ec2 = session.client("ec2") response = asg.describe_auto_scaling_groups(MaxRecords=100) result = response["AutoScalingGroups"] while True: if "NextToken" in response: response = asg.describe_auto_scaling_groups( MaxRecords=100, NextToken=response["NextToken"] ) result = result + response["AutoScalingGroups"] else: break if result and len(result) > 0: for a in result: if "LaunchTemplate" in a: lt_id = a["LaunchTemplate"]["LaunchTemplateId"] lt_version = a["LaunchTemplate"]["Version"] lt_version_result = ec2.describe_launch_template_versions( LaunchTemplateId=lt_id, Versions=[lt_version] ) if ( lt_version_result and "LaunchTemplateVersions" in lt_version_result and len(lt_version_result["LaunchTemplateVersions"]) == 1 ): a["ImageId"] = lt_version_result["LaunchTemplateVersions"][0][ "LaunchTemplateData" ]["ImageId"] return result else: return [] def describe_region_asg_instances(profile) -> list: """ Return all instances in ASG groups in region. """ session = get_boto_session(profile) asg = session.client("autoscaling") response = asg.describe_auto_scaling_instances(MaxRecords=50) result = response["AutoScalingInstances"] while True: if "NextToken" in response: response = asg.describe_auto_scaling_groups( MaxRecords=50, NextToken=response["NextToken"] ) result = result + response["AutoScalingInstances"] else: break return result def _describe_region_eoigw(profile) -> list: """ VPC toolkits include: EOIGW """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_egress_only_internet_gateways(MaxResults=100) result = response["EgressOnlyInternetGateways"] while True: if "NextToken" in response: response = ec2.describe_egress_only_internet_gateways( MaxResults=100, NextToken=response["NextToken"] ) result = result + response["EgressOnlyInternetGateways"] else: break return result def describe_az_ebs(az_name, profile) -> list: """ Returns all EBS in AZ """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_volumes( Filters=[ { "Name": "availability-zone", "Values": [ az_name, ], }, ], MaxResults=100, ) result = response["Volumes"] while True: if "NextToken" in response: response = ec2.describe_volumes( Filters=[ { "Name": "availability-zone", "Values": [ az_name, ], }, ], MaxResults=100, NextToken=response["NextToken"], ) result = result + response["Volumes"] else: break return result def describe_vpc_sg(vpc_id, profile) -> list: """ VPC toolkits include: SG """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_security_groups( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, ) result = response["SecurityGroups"] while True: if "NextToken" in response: response = ec2.describe_security_groups( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, NextToken=response["NextToken"], ) result = result + response["SecurityGroups"] else: break return result def describe_vpc_nacl(vpc_id, profile) -> list: """ VPC toolkits include: NACL """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_network_acls( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, ) result = response["NetworkAcls"] while True: if "NextToken" in response: response = ec2.describe_network_acls( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, NextToken=response["NextToken"], ) result = result + response["NetworkAcls"] else: break return result def describe_vpc_rt(vpc_id, profile) -> list: """ VPC toolkits include: RT """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_route_tables( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, ) result = response["RouteTables"] while True: if "NextToken" in response: response = ec2.describe_route_tables( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, NextToken=response["NextToken"], ) result = result + response["RouteTables"] else: break return result def describe_vpc_by_id(vpc_id, profile) -> dict: try: session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_vpcs(VpcIds=[vpc_id]) result = response["Vpcs"][0] return result except ClientError as e: print(e) def describe_vpc_igw(vpc_id, profile) -> list: """ VPC toolkits include: IGW """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_internet_gateways( Filters=[ { "Name": "attachment.vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, ) result = response["InternetGateways"] while True: if "NextToken" in response: response = ec2.describe_internet_gateways( Filters=[ { "Name": "attachment.vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, StartingToken=response["NextToken"], ) result = result + response["InternetGateways"] else: break return result def describe_vpc_eoigw(vpc_id, profile) -> list: all_eoigws = _describe_region_eoigw(profile) result = [eoigw for eoigw in all_eoigws if eoigw["Attachments"]["VpcId"] == vpc_id] return result def describe_vpc_vpgw(vpc_id, profile) -> list: """ VPC toolkits include: Virtual Private GW """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_vpn_gateways( Filters=[ { "Name": "attachment.vpc-id", "Values": [ vpc_id, ], }, ] ) result = response["VpnGateways"] return result def describe_vpc_endpoint(vpc_id, profile) -> list: """ VPC toolkits include: VPC endpoint """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_vpc_endpoints( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, ) result = response["VpcEndpoints"] while True: if "NextToken" in response: response = ec2.describe_vpc_endpoints( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, NextToken=response["NextToken"], ) result = result + response["VpcEndpoints"] else: break return result def describe_vpc_subnet(vpc_id, profile) -> list: """ Subnet in VPC """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_subnets( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, ) result = response["Subnets"] while True: if "NextToken" in response: response = ec2.describe_subnets( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, ], MaxResults=100, NextToken=response["NextToken"], ) result = result + response["Subnets"] else: break return result def describe_vpc_running_ec2(vpc_id, profile) -> list: """ EC2 in vpc """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_instances( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, { "Name": "instance-state-name", "Values": [ "running", ], }, ], MaxResults=100, ) instances = [] for r in response["Reservations"]: instances += r["Instances"] while True: if "NextToken" in response: response = ec2.describe_instances( Filters=[ { "Name": "vpc-id", "Values": [ vpc_id, ], }, { "Name": "instance-state-name", "Values": [ "running", ], }, ], MaxResults=100, NextToken=response["NextToken"], ) for r in response["Reservations"]: instances += r["Instances"] else: break return instances def describe_subnet_eni(subnet_id, profile) -> list: """ ENIs in subnet """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_network_interfaces( Filters=[ { "Name": "subnet-id", "Values": [ subnet_id, ], }, { "Name": "status", "Values": [ "in-use", ], }, ], MaxResults=100, ) result = response["NetworkInterfaces"] while True: if "NextToken" in response: response = ec2.describe_network_interfaces( Filters=[ { "Name": "subnet-id", "Values": [ subnet_id, ], }, ], MaxResults=100, NextToken=response["NextToken"], ) result += response["NetworkInterfaces"] else: break return result def describe_subnet_ngw(subnet_id, profile) -> list: """ NGW in subnet """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_nat_gateways( Filters=[ { "Name": "subnet-id", "Values": [ subnet_id, ], }, ], MaxResults=100, ) result = response["NatGateways"] while True: if "NextToken" in response: response = ec2.describe_nat_gateways( Filters=[ { "Name": "subnet-id", "Values": [ subnet_id, ], }, ], MaxResults=100, NextToken=response["NextToken"], ) result += response["NatGateways"] else: break return result def describe_subnet_running_ec2(subnet_id, profile) -> list: """ EC2 in subnet """ session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_instances( Filters=[ { "Name": "network-interface.subnet-id", "Values": [ subnet_id, ], }, { "Name": "instance-state-name", "Values": [ "running", ], }, ], MaxResults=100, ) instances = [] for r in response["Reservations"]: instances += r["Instances"] while True: if "NextToken" in response: response = ec2.describe_instances( Filters=[ { "Name": "network-interface.subnet-id", "Values": [ subnet_id, ], }, { "Name": "instance-state-name", "Values": [ "running", ], }, ], MaxResults=100, NextToken=response["NextToken"], ) for r in response["Reservations"]: instances += r["Instances"] else: break return instances # A litte bit complex describers def decsribe_vpcs_from_nametag_prefix(name_prefix, profile) -> list: session = get_boto_session(profile) ec2 = session.client("ec2") response = ec2.describe_vpcs( Filters=[ { "Name": "tag:Name", "Values": [ name_prefix + "*", ], }, ], MaxResults=100, ) result = response["Vpcs"] while True: if "NextToken" in response: response = ec2.describe_vpcs( Filters=[ { "Name": "Name", "Values": [ name_prefix + "*", ], }, ], MaxResults=100, NextToken=response["NextToken"], ) result = result + response["Vpcs"] else: break return result def extract_sg_name_from_sgs_list(sg_id, sgs_list) -> str: if sgs_list: for sg in sgs_list: if sg["GroupId"] == sg_id: return get_nametag_value(sg, "GroupName", sg_id) else: return None def extract_efsfs_name_from_efsfs_list(fs_id, efs_list) -> str: if efs_list: for efs in efs_list: if efs["FileSystemId"] == fs_id: return get_nametag_value(efs, None, fs_id) else: return None def extract_ngw_name_from_ngw_list(ngw_id, ngw_list) -> str: if ngw_list: for ngw in ngw_list: if ngw["NatGatewayId"] == ngw_id: return get_nametag_value(ngw, None, ngw_id) else: return None def extract_eni_name_from_eni_list(eni_id, eni_list) -> str: if eni_list: for eni in eni_list: if eni["NetworkInterfaceId"] == eni_id: return get_nametag_value(eni, None, eni_id) else: return None def extract_instance_name_from_instance_list(instance_id, instance_list) -> str: if instance_list: for i in instance_list: if i["InstanceId"] == instance_id: return get_nametag_value(i, None, instance_id) else: return None def generate_sg_inference_diagram(vpc_id: str, profile=None): profile_prefix = "default" if profile != None: profile_prefix = profile vpc = describe_vpc_by_id(vpc_id, profile) vpc_title = get_nametag_value(vpc, None, vpc_id) vpc_running_instances = describe_vpc_running_ec2(vpc_id, profile) security_groups = describe_vpc_sg(vpc_id, profile) sg_to_permissions = map_sg_id_to_permissions(security_groups) instance_to_nics = {} nic_to_sgs = {} for i in vpc_running_instances: nics = [] for nic in i["NetworkInterfaces"]: nics += [nic["NetworkInterfaceId"]] instance_to_nics[i["InstanceId"]] = nics nsgs = [] for nsg in nic["Groups"]: nsgs += [nsg["GroupId"]] nic_to_sgs[nic["NetworkInterfaceId"]] = nsgs if instance_to_nics != None and nic_to_sgs != None: instance_to_sgs = map_instance_id_to_sgs_list(instance_to_nics, nic_to_sgs) instance_id_to_name = {} for i in [*instance_to_sgs]: instance_id_to_name[i] = extract_instance_name_from_instance_list( i, vpc_running_instances ) else: instance_to_sgs = None # Security groups inference diagram_name_1 = "./stage/" + get_name_by_datetime( profile_prefix + "-" + vpc_title + "-" + "sg-outbound" ) # Outbound with Diagram( "Outbound Traffic Map for Security Groups in " + vpc_title, direction="TB", filename=diagram_name_1, curvestyle="curved", graph_attr={"fontsize": "40"}, edge_attr={"fontsize": "2"}, node_attr={"fontsize": "9"}, ): # print(sg_to_permissions) # Build nodes node_sg_list = set() node_cidr_normal_list = set() node_cidr_reserved_list = set() node_public_list = set() for sg_id in [*sg_to_permissions]: sg_name = extract_sg_name_from_sgs_list(sg_id, security_groups) if sg_name: node_sg_list.add(sg_name) else: node_sg_list.add(sg_id) # All other nodes for l in sg_to_permissions[sg_id]["in"]: for n in l.split(": ")[1].replace(", Protocol", "").split(", "): if "/" in n: seg = n.split(".") if n == "0.0.0.0/0": node_public_list.add(n) elif ( (seg[0] == "192" and seg[1] == "168") or (seg[0] == "10") or ( seg[0] == "172" and int(seg[1]) > 15 and int(seg[1]) < 32 ) ): node_cidr_reserved_list.add(n) else: node_cidr_normal_list.add(n) else: node_sg_list.add(n) for l in sg_to_permissions[sg_id]["out"]: for n in l.split(": ")[1].replace(", Protocol", "").split(", "): if "/" in n: seg = n.split(".") if n == "0.0.0.0/0": node_public_list.add(n) elif ( (seg[0] == "192" and seg[1] == "168") or (seg[0] == "10") or ( seg[0] == "172" and int(seg[1]) > 15 and int(seg[1]) < 32 ) ): node_cidr_reserved_list.add(n) else: node_cidr_normal_list.add(n) else: node_sg_list.add(n) node_full_dict = {} if instance_to_sgs: with Cluster( "EC2 Instance By Name (Note: multiple instances might have the same name)", direction="TB", ): node_instance_dict = { n: EC2(n) for n in set(instance_id_to_name.values()) } with Cluster("Security Group", direction="RL"): node_sg_dict = {n: Endpoint(n) for n in node_sg_list} node_full_dict.update(node_sg_dict) with Cluster("CIDR", direction="TB"): node_cidr_normal_dict = {n: Server(n) for n in node_cidr_normal_list} node_full_dict.update(node_cidr_normal_dict) with Cluster("CIDR Reserved", direction="LR"): node_cidr_reserved_dict = {n: Server(n) for n in node_cidr_reserved_list} node_full_dict.update(node_cidr_reserved_dict) with Cluster("Public", direction="LR"): node_public_dict = {n: Internet(n) for n in node_public_list} node_full_dict.update(node_public_dict) # Draw lines if instance_to_sgs: for i in [*instance_to_sgs]: for sgid in instance_to_sgs[i]: instance_name_to_sg_cache = None for sgid in instance_to_sgs[i]: instance_name_to_sg = sgid + instance_id_to_name[i] if ( instance_name_to_sg_cache != instance_name_to_sg ): # Avoid redrawing the line between the SG and the instances with the same Name tag. ( node_instance_dict[instance_id_to_name[i]] >> Edge(color="purple", style="dotted") >> node_sg_dict[ extract_sg_name_from_sgs_list(sgid, security_groups) ] ) instance_name_to_sg_cache = instance_name_to_sg for sg_id in [*sg_to_permissions]: for l in sg_to_permissions[sg_id]["out"]: for target in l.split(": ")[1].replace(", Protocol", "").split(", "): label = ( l.split(":")[-1] .replace("tcp-ports", "") .replace("icmp-ports", "icmp") .replace("(", "") .replace(")", "") .strip() ) sg_name = extract_sg_name_from_sgs_list(sg_id, security_groups) if sg_name: k = sg_name else: k = sg_id if target in node_public_list: ( node_full_dict[k] >> Edge(color="darkred", label=label) >> node_full_dict[target] ) elif target in node_cidr_reserved_list: ( node_full_dict[k] >> Edge(color="darkgreen", label=label) >> node_full_dict[target] ) elif target in node_cidr_normal_list: ( node_full_dict[k] >> Edge(color="brown", label=label) >> node_full_dict[target] ) elif target in node_sg_list: ( node_full_dict[k] >> Edge(color="darkgrey", label=label) >> node_full_dict[target] ) diagram_name_2 = "./stage/" + get_name_by_datetime( profile_prefix + "-" + vpc_title + "-" + "sg-inbound" ) # Inbound with Diagram( "Inbound Traffic Map for Security Groups in " + vpc_title, direction="TB", filename=diagram_name_2, curvestyle="curved", graph_attr={"fontsize": "40"}, edge_attr={"fontsize": "2"}, node_attr={"fontsize": "9"}, ): # print(sg_to_permissions) # Build nodes node_sg_list = set() node_cidr_normal_list = set() node_cidr_reserved_list = set() node_public_list = set() for sg_id in [*sg_to_permissions]: sg_name = extract_sg_name_from_sgs_list(sg_id, security_groups) if sg_name: node_sg_list.add(sg_name) else: node_sg_list.add(sg_id) # All other nodes for l in sg_to_permissions[sg_id]["in"]: for n in l.split(": ")[1].replace(", Protocol", "").split(", "): if "/" in n: seg = n.split(".") if n == "0.0.0.0/0": node_public_list.add(n) elif ( (seg[0] == "192" and seg[1] == "168") or (seg[0] == "10") or ( seg[0] == "172" and int(seg[1]) > 15 and int(seg[1]) < 32 ) ): node_cidr_reserved_list.add(n) else: node_cidr_normal_list.add(n) else: node_sg_list.add(n) for l in sg_to_permissions[sg_id]["out"]: for n in l.split(": ")[1].replace(", Protocol", "").split(", "): if "/" in n: seg = n.split(".") if n == "0.0.0.0/0": node_public_list.add(n) elif ( (seg[0] == "192" and seg[1] == "168") or (seg[0] == "10") or ( seg[0] == "172" and int(seg[1]) > 15 and int(seg[1]) < 32 ) ): node_cidr_reserved_list.add(n) else: node_cidr_normal_list.add(n) else: node_sg_list.add(n) node_full_dict = {} if instance_to_sgs: with Cluster( "EC2 Instance By Name (Note: multiple instances might have the same name)", direction="TB", ): node_instance_dict = { n: EC2(n) for n in set(instance_id_to_name.values()) } with Cluster("Security Group", direction="RL"): node_sg_dict = {n: Endpoint(n) for n in node_sg_list} node_full_dict.update(node_sg_dict) with Cluster("CIDR", direction="TB"): node_cidr_normal_dict = {n: Server(n) for n in node_cidr_normal_list} node_full_dict.update(node_cidr_normal_dict) with Cluster("CIDR Reserved", direction="LR"): node_cidr_reserved_dict = {n: Server(n) for n in node_cidr_reserved_list} node_full_dict.update(node_cidr_reserved_dict) with Cluster("Public", direction="LR"): node_public_dict = {n: Internet(n) for n in node_public_list} node_full_dict.update(node_public_dict) # Draw lines if instance_to_sgs: for i in [*instance_to_sgs]: instance_name_to_sg_cache = None for sgid in instance_to_sgs[i]: instance_name_to_sg = sgid + instance_id_to_name[i] if ( instance_name_to_sg_cache != instance_name_to_sg ): # Avoid redrawing the line between the SG and the instances with the same Name tag. ( node_instance_dict[instance_id_to_name[i]] >> Edge(color="purple", style="dotted") >> node_sg_dict[ extract_sg_name_from_sgs_list(sgid, security_groups) ] ) instance_name_to_sg_cache = instance_name_to_sg # print(node_full_dict) for sg_id in [*sg_to_permissions]: for l in sg_to_permissions[sg_id]["in"]: for source in l.split(": ")[1].replace(", Protocol", "").split(", "): seg = source.split(".") label = ( l.split(":")[-1] .replace("tcp-ports", "") .replace("icmp-ports", "icmp") .replace("(", "") .replace(")", "") .strip() ) sg_name = extract_sg_name_from_sgs_list(sg_id, security_groups) if sg_name: k = sg_name else: k = sg_id if source in node_public_list: ( node_full_dict[source] >> Edge(color="darkred", label=label) >> node_full_dict[k] ) elif source in node_cidr_reserved_list: ( node_full_dict[source] >> Edge(color="darkgreen", label=label) >> node_full_dict[k] ) elif source in node_cidr_normal_list: ( node_full_dict[source] >> Edge(color="brown", label=label) >> node_full_dict[k] ) elif source in node_sg_list: ( node_full_dict[source] >> Edge(color="darkgrey", label=label) >> node_full_dict[k] )