added pydot dependency, added visualization

principalmapper/__main__.py

@@ -15,6 +15,7 @@
 from principalmapper.util import botocore_tools
 from principalmapper.util.debug_print import dprint
 from principalmapper.util.storage import get_storage_root
+from principalmapper.visualizing import graph_writer
 
 
 def main():
@@ -139,7 +140,7 @@ def main():
     visualizationparser.add_argument(
         '--filetype',
         default='svg',
-        choices=['svg', 'png'],
+        choices=['svg', 'png', 'dot'],
         help='The (lowercase) filetype to output the image as.'
     )
 
@@ -267,7 +268,16 @@ def handle_repl(parsed_args):
 
 def handle_visualization(parsed_args):
     """Processes the arguments for the visualization subcommand and executes related tasks"""
-    raise NotImplementedError('visualize subcommand is not ready for use')  # TODO: visualization functionality
+    # get Graph to draw/write
+    if parsed_args.account is None:
+        session = botocore_tools.get_session(parsed_args.profile)
+    else:
+        session = None
+    graph = principalmapper.graphing.graph_actions.get_existing_graph(session, parsed_args.account, parsed_args.debug)
+
+    # create file
+    filepath = './{}.{}'.format(graph.metadata['account_id'], parsed_args.filetype)
+    graph_writer.handle_request(graph, filepath, parsed_args.filetype)
 
 
 def handle_analysis(parsed_args):

principalmapper/visualizing/graph_writer.py

@@ -0,0 +1,38 @@
+"""Code to write Graph data to various output formats."""
+
+import pydot
+
+from principalmapper.common.graphs import Graph
+from principalmapper.querying.presets.privesc import can_privesc
+
+
+def handle_request(graph: Graph, path: str, format: str) -> None:
+    """Meat of the graph_writer.py module, writes graph data in a given file-format to the given path."""
+    # Load graph data into pydot
+    pydg = pydot.Dot(
+        graph_type='digraph',
+        graph_name='Principal Mapper Visualization: {}'.format(graph.metadata['account_id']),
+        overlap='scale',
+        layout='neato',
+        concentrate='true',
+        splines='true'
+    )
+    pyd_nd = {}
+
+    for node in graph.nodes:
+        if node.is_admin:
+            color = '#BFEFFF'
+        elif can_privesc(graph, node)[0]:
+            color = '#FADBD8'
+        else:
+            color = 'white'
+
+        pyd_nd[node] = pydot.Node(node.searchable_name(), style='filled', fillcolor=color, shape='box')
+        pydg.add_node(pyd_nd[node])
+
+    for edge in graph.edges:
+        if not edge.source.is_admin:
+            pydg.add_edge(pydot.Edge(pyd_nd[edge.source], pyd_nd[edge.destination]))
+
+    # and draw
+    pydg.write(path, format=format)

setup.py

@@ -16,7 +16,7 @@
     scripts=['pmapper.py'],
     packages=['principalmapper'],
     package_data={},
-    install_requires=['botocore', 'packaging', 'python-dateutil'],
+    install_requires=['botocore', 'packaging', 'python-dateutil', 'pydot'],
     entry_points={
         'console_scripts': [
             'pmapper = principalmapper.__main__:main'