feat: add orgs and scp visualization capability (#25)

* feat: add orgs and scp visualization capability

* docs: add visualization and update bytes to character limit

Author: dgwhited

README.md

@@ -5,16 +5,18 @@
 
 This project provides a Python module to aid in Service Control Policy (SCP) management in AWS accounts.
 
-SCPs have a current limit of 5 total per entity, and a size limit on each of 5120 bytes. This tool will merge selected SCPs into the fewest amount of policies, and optionally remove whitespace characters as they count toward the byte limit.
+SCPs have a current limit of 5 total per entity, and a size limit on each of 5120 characters. This tool will merge selected SCPs into the fewest amount of policies, and optionally remove whitespace characters as they count toward the character limit.
 
 
 ```mermaid
   stateDiagram-v2
       [SCPTool] --> Validate
       [SCPTool] --> Merge
+      [SCPTool] --> Visualize
       Merge --> Validate
       Validate --> [*]
       Merge --> [*]
+      Visualize --> [*]
 ```
 ## Using SCPkit
 SCPkit can be installed from PyPI
@@ -38,23 +40,36 @@ Optional validation with output locally:
 scpkit merge --sourcefiles /path/to/scps --outdir /path/to/directory --validate-after-merge --profile yourawsprofile
 ```
 
+### Creating a visualization of an AWS Organization, OUs, Accounts, and SCPs
+Creating this visualization requires you be authenticated with either the Org management account, or a delegated administrator. See the [AWS Documentation](https://docs.aws.amazon.com/organizations/latest/userguide/orgs_delegate_policies_example_view_accts_orgs.html) page for more info on delegating Organizations.
+
+This will output a graph pdf and graphviz data file in the specified directory (or local directory, if outdir is not specified.)
+
+```
+scpkit visualize --profile yourawsprofile --outdir ./org-graph
+```
+Accounts are presented as ellipses, organizational units are rectangles, and SCPs are trapezoids.
+
+![Visualization of an Organization](./visualize-org.png)
+
 The full CLI is documented through docopt
 ```
-SCPkit
+"""SCPkit
 Usage:
-    main.py (validate | merge) [--sourcefiles sourcefiles] [--profile profile] [ --outdir outdir] [--validate-after-merge] [--readable]
+    main.py (validate | merge | visualize) [--sourcefiles sourcefiles] [--profile profile] [ --outdir outdir] [--validate-after-merge] [--readable] [--console]
 
 Options:
     -h --help                   Show this screen.
     --version                   Show version.
-    --sourcefiles sourcefiles   Directory path to SCP files in json format
+    --sourcefiles sourcefiles   Directory path to SCP files in json format or a single SCP file
     --outdir outdir             Directory to write new SCP files [Default: ./]
     --profile profile           AWS profile name
     --validate-after-merge      Validate the policies after merging them
     --readable                  Leave indentation and some whitespace to make the SCPs readable
+    --console                   Adds Log to console
+"""
 ```
 
-
 ## Local development
 From the root of the folder:
 ```

requirements.txt

@@ -1,2 +1,3 @@
-boto3
-docopt
+boto3>=1.28.66
+docopt>=0.6.2
+graphviz>=0.20.1

scpkit/main.py

@@ -1,6 +1,6 @@
 """SCPkit
 Usage:
-    main.py (validate | merge) [--sourcefiles sourcefiles] [--profile profile] [ --outdir outdir] [--validate-after-merge] [--readable] [--console]
+    main.py (validate | merge | visualize) [--sourcefiles sourcefiles] [--profile profile] [ --outdir outdir] [--validate-after-merge] [--readable] [--console]
 
 Options:
     -h --help                   Show this screen.
@@ -16,13 +16,17 @@
 from .src.validate import validate_policies
 from .src.merge import scp_merge
 from .src.util import get_files_in_dir
+from .src.visualize import visualize_policies
 
 def main():
     arguments = {
         k.lstrip('-'): v for k, v in docopt(__doc__).items()
     }
 
-    arguments['scps'] = get_files_in_dir(arguments["sourcefiles"])
+    if arguments.get("visualize"):
+        visualize_policies(arguments['profile'], arguments['outdir'])
+    else:
+        arguments['scps'] = get_files_in_dir(arguments["sourcefiles"])
 
     if arguments.get("merge"):
         scp_merge(**arguments)

scpkit/src/merge.py

@@ -30,12 +30,12 @@ def merge_json(json_blobs):
 
 
 def make_policies(content, readable, max_size: int = 5120):
-    """Combines the policies in order, counts the bytes, and starts a new file when it goes over the limit.
+    """Combines the policies in order, counts the characters, and starts a new file when it goes over the limit.
         Theres probably a better way to do this with permutations, but that could also be resource intensive.
 
     Args:
         content (list): List of Sid dictionaries (in order of smallest to largest preferred)
-        max_size (int, optional): Max byte count. Defaults to 5120.
+        max_size (int, optional): Max character count. Defaults to 5120.
     Returns:
         list: List of condensed SCP documents.
     """

scpkit/src/util.py

@@ -1,4 +1,5 @@
 import json
+import boto3
 from pathlib import Path
 from .model import SCP
 
@@ -112,4 +113,48 @@ def make_actions_and_resources_lists(content):
         # no such thing as NotResource in SCPs - https://docs.aws.amazon.com/organizations/latest/userguide/orgs_manage_policies_scps_syntax.html#scp-elements-table
         if type(sid.get("Resource")) is not list:
             sid["Resource"] = [sid.get("Resource")]
-    return content
+    return content
+
+
+def create_session(profile=None):
+    """Creates a boto session
+
+    Args:
+        profile (string): AWS profile name
+
+    Returns:
+        [object]: Authenticated Boto3 session
+    """
+    if profile:
+        return boto3.Session(profile_name=profile)
+    else:
+        return boto3.Session()
+
+
+def create_client(session, service):
+    """Creates a service client from a boto session
+
+    Args:
+        session (object): Authenicated boto3 session
+        service (string): service name to create the client for
+
+    Returns:
+        [object]: client session for specific aws service (eg. accessanalyzer)
+    """
+    return session.client(service)
+
+
+def paginate(service, method, **method_args):
+    """Paginates through the results of a method.
+
+    Args:
+        service (boto3.client): The AWS service client.
+        method (str): The name of the method to paginate.
+        method_args (dict): The arguments to pass to the method.
+
+    Returns:
+        list: A list of paginated results.
+    """
+    paginator = service.get_paginator(method)
+    results = paginator.paginate(**method_args)
+    return results

scpkit/src/validate.py

@@ -1,31 +1,4 @@
-import boto3
-
-def create_session(profile=None):
-    """Creates a boto session
-
-    Args:
-        profile (string): AWS profile name
-
-    Returns:
-        [object]: Authenticated Boto3 session
-    """
-    if profile:
-        return boto3.Session(profile_name=profile)    
-    else:
-        return boto3.Session()
-
-
-def create_client(session, service):
-    """Creates a service client from a boto session
-
-    Args:
-        session (object): Authenicated boto3 session
-        service (string): service name to create the client for
-
-    Returns:
-        [object]: client session for specific aws service (eg. accessanalyzer)
-    """
-    return session.client(service)
+from .util import create_session, create_client
 
 
 def validate_policies(scps, profile, outdir=None, console=False):

scpkit/src/visualize.py

@@ -0,0 +1,118 @@
+from graphviz import Digraph
+from .util import create_session, paginate
+
+
+def add_child_nodes(ou_id, org_client, graph):
+    """Adds child nodes to the graph.
+
+    Args:
+        ou_id (str): The ID of the organizational unit.
+        org_client (boto3.client): The AWS Organizations client.
+        graph (Digraph): The Graphviz Digraph object.
+
+    """
+    accounts = list_children(org_client, ou_id, 'ACCOUNT')
+    ous = list_children(org_client, ou_id, 'ORGANIZATIONAL_UNIT')
+
+    children = accounts + ous
+
+    if children:
+        for child in children:
+            child_id = child['Id']
+            child_type = child['Type']
+
+            if child_type == 'ACCOUNT':
+                account = org_client.describe_account(AccountId=child_id).get('Account')
+                account_name = account.get('Name')
+                account_id = account.get('Id')
+
+                graph.node(child_id, label=account_name, shape='ellipse')
+                graph.edge(ou_id, child_id)
+
+                policies = get_policies_for_entity(account_id, org_client)
+
+                add_policies_to_graph(graph, child_id, policies=policies)
+
+            # Get the name of the child (OU or Account)
+            if child_type == 'ORGANIZATIONAL_UNIT':
+                current_ou = org_client.describe_organizational_unit(OrganizationalUnitId=child_id).get('OrganizationalUnit')
+                ou_name = current_ou.get('Name')
+                current_ou_id = current_ou.get('Id')
+
+                graph.node(child_id, label=ou_name, shape='box')
+                graph.edge(ou_id, child_id)
+
+                policies = get_policies_for_entity(current_ou_id, org_client)
+                add_policies_to_graph(graph, child_id, policies=policies)
+
+                add_child_nodes(child_id, org_client, graph)
+
+
+def list_children(org_client, parent_id, child_type):
+    """Lists the children of a parent entity.
+
+    Args:
+        org_client (boto3.client): The AWS Organizations client.
+        parent_id (str): The ID of the parent entity.
+        child_type (str): The type of the child entities to list.
+
+    Returns:
+        list: A list of child entities.
+    """
+    all_children = paginate(org_client, 'list_children', ParentId=parent_id, ChildType=child_type)
+    children = [ child for page in all_children for child in page.get('Children')]
+    return children
+
+
+def get_policies_for_entity(entity_id, org_client, filter='SERVICE_CONTROL_POLICY'):
+    """Gets the policies associated with an entity.
+
+    Args:
+        entity_id (str): The ID of the entity.
+        org_client (boto3.client): The AWS Organizations client.
+        filter (str): The filter to apply when retrieving policies.
+
+    Returns:
+        list: A list of policies associated with the entity.
+    """
+    policies = org_client.list_policies_for_target(
+        TargetId=entity_id,
+        Filter=filter
+    )
+    return policies.get('Policies')
+
+
+def add_policies_to_graph(graph, entity_id, policies=None):
+    """Adds policies to the graph.
+
+    Args:
+        graph (Digraph): The Graphviz Digraph object.
+        entity_id (str): The ID of the entity.
+        policies (list): A list of policies to add to the graph.
+    """
+    if policies:
+        for policy in policies:
+            policy_name = policy.get('Name')
+            graph.node(policy_name, label=policy_name, shape='trapezium')
+            graph.edge(entity_id, policy_name)
+
+
+def visualize_policies(profile, outdir):
+    session = create_session(profile)
+    org_client = session.client('organizations')
+
+    # Initialize a Graphviz Digraph object
+    graph = Digraph('AWS_Organizations', graph_attr={'rankdir':'LR'})
+
+    # Get the root information
+    root_id = org_client.list_roots()['Roots'][0]['Id']
+    graph.node(root_id, label="Root", shape='box')
+    get_policies_for_entity(root_id, org_client)
+    add_policies_to_graph(graph, root_id)
+
+    # Start building the tree
+    add_child_nodes(root_id, org_client, graph)
+
+    # Output the graphical tree hierarchy to a file
+    graph.render(directory=outdir, filename='aws_org_tree', view=True)
+