Added integration with Diagrams project #16

.gitignore

@@ -32,6 +32,7 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+assets/diagrams/
 
 # PyInstaller
 #  Usually these files are written by a python script from a template

aws-network-discovery.py

@@ -39,6 +39,7 @@
 __version__ = "0.7.0"
 
 AVAILABLE_LANGUAGES = ['en_US','pt_BR']
+DIAGRAMS_OPTIONS = ['True','False']
 
 def show_options(args="sys.argv[1:]"):
     parser = argparse.ArgumentParser()
@@ -66,6 +67,13 @@ def show_options(args="sys.argv[1:]"):
                         required=False,
                         help="available languages: pt_BR, en_US"
                         )
+    parser.add_argument(
+                        "-d",
+                        "--diagram",
+                        required=False,
+                        help="print diagram with resources (need Graphviz installed). Use options \"True\" to " \
+                             "view image or \"False\" to save image to disk. Default True"
+                        )
     args = parser.parse_args()
 
     return args
@@ -85,14 +93,22 @@ def main():
     else:
         language = args.language
 
+    """ Diagram check """
+    if args.diagram is not None and args.diagram not in DIAGRAMS_OPTIONS:
+        diagram = "True"
+    else:
+        diagram = args.diagram
 
     """ defining default language to show messages """
     defaultlanguage = gettext.translation('messages', localedir='locales', languages=[language])
     defaultlanguage.install()
     _ = defaultlanguage.gettext 
 
 
-    vpc = Vpc(vpc_id=args.vpc_id, region_name=args.region_name, profile_name=args.profile_name)
+    vpc = Vpc(vpc_id=args.vpc_id, 
+              region_name=args.region_name, 
+              profile_name=args.profile_name,
+              diagram=diagram)
     vpc.run()
 
 

commands/vpc.py

@@ -4,10 +4,11 @@
 
 class Vpc(object):
 
-    def __init__(self, vpc_id, region_name, profile_name):
+    def __init__(self, vpc_id, region_name, profile_name, diagram):
         self.vpc_id = vpc_id
         self.region_name = region_name
         self.profile_name = profile_name
+        self.diagram = diagram
 
     def run(self):
 
@@ -29,10 +30,12 @@ def run(self):
             vpcs = client.describe_vpcs()
             for data in vpcs['Vpcs']:
                 """ init class awscommands """
-                awscommands = AwsCommands(VpcOptions(session=session, vpc_id=data['VpcId'], region_name=region_name)).run()
+                awscommands = AwsCommands(VpcOptions(session=session, vpc_id=data['VpcId'], region_name=region_name),
+                                                     diagram=self.diagram).run()
         else:
             """ init class awscommands """
-            awscommands = AwsCommands(VpcOptions(session=session, vpc_id=self.vpc_id, region_name=region_name)).run()
+            awscommands = AwsCommands(VpcOptions(session=session, vpc_id=self.vpc_id, region_name=region_name),
+                                                 diagram=self.diagram).run()
 
 
 

shared/awscommands.py

@@ -2,15 +2,17 @@
 from shared.internal.security import IAM
 from shared.internal.network import VPC
 from shared.report import Report
+from shared.diagram import _Diagram
 import importlib, inspect
 import os
 
 PATH_CHECKS = "shared/internal"
 
 class AwsCommands(object):
 
-    def __init__(self, vpc_options: VpcOptions):
+    def __init__(self, vpc_options: VpcOptions, diagram):
         self.vpc_options = vpc_options
+        self.diagram = diagram
 
     def run(self):
 
@@ -49,9 +51,10 @@ def run(self):
         Report(resources=resources_check).generalReport()
 
         """ 
-        TODO: Generate diagrams... future...
+        Diagram integration
         """
-        #....diagrams(checks)....
+        if self.diagram is not None:
+            _Diagram(vpc_id=self.vpc_options.vpc_id, diagram=self.diagram, resources=resources_check).generateDiagram()
 
         """
         TODO: Export in csv/json/yaml/tf... future...

shared/common.py

@@ -34,6 +34,7 @@ class Resource(NamedTuple):
     name: str
     type: str
     details: str
+    group: str
 
 
 

shared/diagram.py

@@ -0,0 +1,96 @@
+from shared.common import *
+from shared.error_handler import exception
+import os
+from diagrams import Cluster, Diagram
+""" Importing all AWS nodes """
+from diagrams.aws.analytics import *
+from diagrams.aws.compute import *
+from diagrams.aws.database import *
+from diagrams.aws.devtools import *
+from diagrams.aws.engagement import *
+from diagrams.aws.integration import *
+from diagrams.aws.iot import *
+from diagrams.aws.management import *
+from diagrams.aws.media import *
+from diagrams.aws.migration import *
+from diagrams.aws.ml import *
+from diagrams.aws.network import *
+from diagrams.aws.security import *
+from diagrams.aws.storage import *
+
+
+PATH_DIAGRAM_OUTPUT = "./assets/diagrams/"
+
+class Mapsources:
+
+    """ Class to mapping type resource from Terraform to Diagram Nodes """
+    mapresources = {"aws_lambda_function": "Lambda", "aws_emr_cluster": "EMRCluster",
+                    "aws_elasticsearch_domain": "ES", "aws_msk_cluster": "ManagedStreamingForKafka",
+                    "aws_sqs_queue_policy": "SQS", "aws_instance": "EC2",
+                    "aws_eks_cluster": "EKS", "aws_autoscaling_group": "AutoScaling",
+                    "aws_ecs_cluster": "ECS", "aws_db_instance": "RDS",
+                    "aws_elasticache_cluster": "ElastiCache", "aws_docdb_cluster": "DocumentDB",
+                    "aws_internet_gateway": "InternetGateway", "aws_nat_gateway": "NATGateway",
+                    "aws_elb_classic": "ELB", "aws_elb": "ELB",
+                    "aws_route_table": "RouteTable", "aws_subnet": "PublicSubnet",
+                    "aws_network_acl": "Nacl", "aws_vpc_peering_connection": "VPCPeering",
+                    "aws_vpc_endpoint_gateway": "Endpoint", "aws_iam_policy": "IAM",
+                    "aws_efs_file_system": "EFS", "aws_s3_bucket_policy": "S3"}
+
+
+class _Diagram(object):
+
+    def __init__ (self, vpc_id, diagram, resources):
+        self.resources = resources
+        self.vpc_id = vpc_id
+        self.diagram = diagram
+
+    @exception
+    def generateDiagram(self):
+
+        #diagram = [[], [], [], [], [], [], [Resource(id='j-HCK3AB8SJ0JK', name='My cluster', type='aws_emr_cluster', details='', group='compute'), Resource(id='j-UU8T5E0AI99S', name='My cluster', type='aws_emr_cluster', details='', group='compute')], [Resource(id='arn:aws:lambda:us-east-1:200984112386:function:cwsyn-meucana-2b5c2c74-cbfc-46b5-8400-db853e98dca3', name='cwsyn-meucana-2b5c2c74-cbfc-46b5-8400-db853e98dca3', type='aws_lambda_function', details='', group='compute'), Resource(id='arn:aws:lambda:us-east-1:200984112386:function:myfnn', name='myfnn', type='aws_lambda_function', details='', group='compute'), Resource(id='arn:aws:lambda:us-east-1:200984112386:function:aaaaaaaaaaaaa', name='aaaaaaaaaaaaa', type='aws_lambda_function', details='', group='compute')], [], [], [], [], [], [], [], [Resource(id='igw-09cc041a3f6e7d38b', name='igw-09cc041a3f6e7d38b', type='aws_internet_gateway', details='', group='network')], [Resource(id='acl-01a0765c917e4e733', name='acl-01a0765c917e4e733', type='aws_network_acl', details='NACL using Subnets subnet-0267b72fb6dead183, subnet-023da38a98c179cce, subnet-0c63113a164c83e0b', group='network')], [], [Resource(id='rtb-0405795a9a7f730e8', name='rtb-0405795a9a7f730e8', type='aws_route_table', details='', group='network'), Resource(id='rtb-0bf969c6f52d3ca40', name='rtb-0bf969c6f52d3ca40', type='aws_route_table', details='', group='network')], [Resource(id='sg-0521cf0ebff65c70d', name='ElasticMapReduce-slave', type='aws_security_group', details='', group='network'), Resource(id='sg-0abf4b1f1a9ecfe9c', name='MEUSG', type='aws_security_group', details='', group='network'), Resource(id='sg-0b21dea57386a36b4', name='ElasticMapReduce-master', type='aws_security_group', details='', group='network'), Resource(id='sg-0f2d79f705ff0d97d', name='default', type='aws_security_group', details='', group='network')], [Resource(id='subnet-0c63113a164c83e0b', name='subnet-0c63113a164c83e0b', type='aws_subnet', details='Subnet using CidrBlock 10.0.2.0/24 and AZ us-east-1e', group='network'), Resource(id='subnet-0267b72fb6dead183', name='subnet-0267b72fb6dead183', type='aws_subnet', details='Subnet using CidrBlock 10.0.3.0/24 and AZ us-east-1b', group='network'), Resource(id='subnet-023da38a98c179cce', name='subnet-023da38a98c179cce', type='aws_subnet', details='Subnet using CidrBlock 10.0.1.0/24 and AZ us-east-1a', group='network')], [], [], [], [], []]
+        """ Check if assets/diagram directory exists """
+        if not os.path.isdir(PATH_DIAGRAM_OUTPUT):
+            print("OI")
+            try:
+                os.mkdir(PATH_DIAGRAM_OUTPUT)
+            except OSError:
+                print ("Creation of the directory %s failed" % PATH_DIAGRAM_OUTPUT)
+            else:
+                print ("Successfully created the directory %s " % PATH_DIAGRAM_OUTPUT)
+
+        """ Ordering Resource list to group resources into cluster """
+        ordered_resources = dict()
+        for alldata in self.resources:
+            if isinstance(alldata, list):
+                for rundata in alldata:
+                    if Mapsources.mapresources.get(rundata.type) is not None:
+                        if rundata.group in ordered_resources:
+                            ordered_resources[rundata.group].append({"id": rundata.id,
+                                                                    "type": rundata.type,
+                                                                    "name": rundata.name,
+                                                                    "details": rundata.details})
+                        else:
+                            ordered_resources[rundata.group] = [{"id": rundata.id,
+                                                                "type": rundata.type,
+                                                                "name": rundata.name,
+                                                                "details": rundata.details}]
+
+        """ Start mounting Cluster """
+        resource_id = list()
+        with Diagram(name="AWS VPC {} Resources".format(self.vpc_id), filename=PATH_DIAGRAM_OUTPUT+self.vpc_id, 
+                     show=self.diagram, direction="TB"):
+
+            """ VPC to represent main resource """
+            _vpc = VPC("VPC {}".format(self.vpc_id))
+
+            """ Iterate resources to draw it """
+            for alldata in ordered_resources:
+                with Cluster(alldata.capitalize() + " resources"):
+                    for rundata in ordered_resources[alldata]:
+                        resource_id.append(eval(Mapsources.mapresources.get(rundata["type"]))(rundata["name"]))
+
+
+            """ Connecting resources and vpc """
+            for resource in resource_id:
+                resource >> _vpc

shared/internal/analytics.py

@@ -39,7 +39,8 @@ def run(self) -> List[Resource]:
                     resources_found.append(Resource(id=elasticsearch_domain['DomainStatus']['DomainId'],
                                                     name=elasticsearch_domain['DomainStatus']['DomainName'],
                                                     type='aws_elasticsearch_domain',
-                                                    details=''))
+                                                    details='',
+                                                    group='analytics'))
 
         return resources_found
 
@@ -81,7 +82,8 @@ def run(self) -> List[Resource]:
                         resources_found.append(Resource(id=data['ClusterArn'],
                                                         name=data['ClusterName'],
                                                         type='aws_msk_cluster',
-                                                        details=''))
+                                                        details='',
+                                                        group='analytics'))
 
                         break
         return resources_found

shared/internal/application.py

@@ -55,6 +55,7 @@ def analyze_queues(self, client, queue):
                     return True, Resource(id=queuearn,
                                            name=queue,
                                            type='aws_sqs_queue_policy',
-                                           details='')
+                                           details='',
+                                           group='application')
 
         return False, None

shared/internal/compute.py

@@ -27,7 +27,8 @@ def run(self) -> List[Resource]:
                     resources_found.append(Resource(id=data['FunctionArn'],
                                           name=data["FunctionName"],
                                           type='aws_lambda_function',
-                                          details=''))
+                                          details='',
+                                          group='compute'))
 
         return resources_found
 
@@ -62,7 +63,8 @@ def run(self) -> List[Resource]:
                             resources_found.append(Resource(id=instances['InstanceId'],
                                                             name=instance_name,
                                                             type='aws_instance',
-                                                            details=''))
+                                                            details='',
+                                                            group='compute'))
 
 
         return resources_found
@@ -94,7 +96,8 @@ def run(self) -> List[Resource]:
                     resources_found.append(Resource(id=cluster['cluster']['arn'],
                                                     name=cluster['cluster']["name"],
                                                     type='aws_eks_cluster',
-                                                    details=''))
+                                                    details='',
+                                                    group='compute'))
 
         return  resources_found
 
@@ -132,7 +135,8 @@ def run(self) -> List[Resource]:
                     resources_found.append(Resource(id=data['Id'],
                                                     name=data['Name'],
                                                     type='aws_emr_cluster',
-                                                    details=''))
+                                                    details='',
+                                                    group='compute'))
 
         return resources_found
 
@@ -171,6 +175,7 @@ def run(self) -> List[Resource]:
                         resources_found.append(Resource(id=data['AutoScalingGroupARN'],
                                                         name=data['AutoScalingGroupName'],
                                                         type='aws_autoscaling_group',
-                                                        details='Using LaunchConfigurationName {0}'.format(data["LaunchConfigurationName"])))
+                                                        details='Using LaunchConfigurationName {0}'.format(data["LaunchConfigurationName"]),
+                                                        group='compute'))
 
         return resources_found

shared/internal/containers.py

@@ -57,7 +57,8 @@ def run(self) -> List[Resource]:
                                         resources_found.append(Resource(id=data['clusterArn'],
                                                                         name=data["clusterName"],
                                                                         type='aws_ecs_cluster',
-                                                                        details=''))
+                                                                        details='',
+                                                                        group='container'))
                             else:
                                 """ EC2 services require container instances, list of them should be fine for now """
                                 pass
@@ -91,7 +92,8 @@ def run(self) -> List[Resource]:
                                         resources_found.append(Resource(id=instance['InstanceId'],
                                                                         name=data["clusterName"],
                                                                         type='aws_ecs_cluster',
-                                                                        details='Instance in EC2 cluster'))
+                                                                        details='Instance in EC2 cluster',
+                                                                        group='container'))
                                     pass
                             pass
                         pass

shared/internal/database.py

@@ -36,7 +36,8 @@ def run(self) -> List[Resource]:
                                                     name=data["DBInstanceIdentifier"],
                                                     type='aws_db_instance',
                                                     details='DBInstance using subnets {} and engine {}'\
-                                                    .format(', '.join(subnet_ids), data["Engine"])))
+                                                    .format(', '.join(subnet_ids), data["Engine"]),
+                                                    group='database'))
 
         return resources_found
 
@@ -74,7 +75,8 @@ def run(self) -> List[Resource]:
                                                     name=data["CacheSubnetGroupName"],
                                                     type='aws_elasticache_cluster',
                                                     details='Elasticache Cluster using subnets {} and engine {}' \
-                                                    .format(', '.join(subnet_ids), data["Engine"])))
+                                                    .format(', '.join(subnet_ids), data["Engine"]),
+                                                    group='database'))
 
         return resources_found
 
@@ -112,6 +114,7 @@ def run(self) -> List[Resource]:
                                                     name=data["DBInstanceIdentifier"],
                                                     type='aws_docdb_cluster',
                                                     details='Documentdb using subnets {} and engine {}'\
-                                                    .format(', '.join(subnet_ids), data["Engine"])))
+                                                    .format(', '.join(subnet_ids), data["Engine"]),
+                                                    group='database'))
 
         return resources_found

shared/internal/management.py

@@ -31,6 +31,7 @@ def run(self) -> List[Resource]:
                         resources_found.append(Resource(id=data['Id'],
                                                         name=data["Name"],
                                                         type='aws_canaries_function',
-                                                        details=''))
+                                                        details='',
+                                                        group='management'))
 
         return resources_found