fixed up AWS Organizations code for non-aws partitions

Author: Erik Steringer

principalmapper/graphing/graph_cli.py

@@ -28,7 +28,7 @@
 from principalmapper.graphing.gathering import get_organizations_data
 from principalmapper.graphing.edge_identification import checker_map
 from principalmapper.querying import query_orgs
-from principalmapper.util import botocore_tools
+from principalmapper.util import botocore_tools, arns
 from principalmapper.util.storage import get_storage_root, get_default_graph_path
 
 
@@ -149,9 +149,13 @@ def process_arguments(parsed_args: Namespace):
                 stsclient = session.create_client('sts')
             caller_identity = stsclient.get_caller_identity()
             caller_account = caller_identity['Account']
+            partition = arns.get_partition(caller_identity['Arn'])
             logger.debug("Caller Identity: {}".format(caller_identity))
 
-            org_tree_search_dir = Path(get_storage_root())
+            if partition == 'aws':
+                org_tree_search_dir = Path(get_storage_root())
+            else:
+                org_tree_search_dir = Path(os.path.join(get_storage_root(), partition))
             org_id_pattern = re.compile(r'/o-\w+')
             for subdir in org_tree_search_dir.iterdir():
                 if org_id_pattern.search(str(subdir)) is not None:

principalmapper/graphing/orgs_cli.py

@@ -29,7 +29,7 @@
 from principalmapper.graphing.gathering import get_organizations_data
 from principalmapper.querying.query_orgs import produce_scp_list
 from principalmapper.util import botocore_tools
-from principalmapper.util.storage import get_storage_root
+from principalmapper.util.storage import get_storage_root, get_default_graph_path
 
 
 logger = logging.getLogger(__name__)
@@ -106,15 +106,20 @@ def process_arguments(parsed_args: Namespace):
         # create the account -> OU path map and apply to all accounts (same as orgs update operation)
         account_ou_map = _map_account_ou_paths(org_tree)
         logger.debug('account_ou_map: {}'.format(account_ou_map))
-        _update_accounts_with_ou_path_map(org_tree.org_id, account_ou_map, get_storage_root())
+        root_path = get_storage_root() if org_tree.partition == 'aws' else os.path.join(get_storage_root(), org_tree.partition)
+        _update_accounts_with_ou_path_map(org_tree.org_id, account_ou_map, root_path)
         logger.info('Updated currently stored Graphs with applicable AWS Organizations data')
 
         # create and cache a list of edges between all the accounts we have data for
         edge_list = []
         graph_objs = []
         for account in org_tree.accounts:
             try:
-                potential_path = os.path.join(get_storage_root(), account)
+                if org_tree.partition != 'aws':
+                    potential_path = get_default_graph_path(f'{org_tree.partition}:{account}')
+                else:
+                    potential_path = get_default_graph_path(account)
+
                 logger.debug('Trying to load a Graph from {}'.format(potential_path))
                 graph_obj = Graph.create_graph_from_local_disk(potential_path)
                 graph_objs.append(graph_obj)
@@ -135,26 +140,34 @@ def process_arguments(parsed_args: Namespace):
         org_tree.edge_list = edge_list
         logger.info('Compiled cross-account edges')
 
-        org_tree.save_organization_to_disk(os.path.join(get_storage_root(), org_tree.org_id))
+        if org_tree.partition != 'aws':
+            org_storage_path = get_default_graph_path(f'{org_tree.partition}:{org_tree.org_id}')
+        else:
+            org_storage_path = get_default_graph_path(org_tree.org_id)
+        org_tree.save_organization_to_disk(org_storage_path)
         logger.info('Stored organization data to disk')
 
     elif parsed_args.picked_orgs_cmd == 'update':
         # pull the existing data from disk
-        org_filepath = os.path.join(get_storage_root(), parsed_args.org)
+        org_filepath = get_default_graph_path(parsed_args.org)
         org_tree = OrganizationTree.create_from_dir(org_filepath)
 
         # create the account -> OU path map and apply to all accounts
         account_ou_map = _map_account_ou_paths(org_tree)
         logger.debug('account_ou_map: {}'.format(account_ou_map))
-        _update_accounts_with_ou_path_map(org_tree.org_id, account_ou_map, get_storage_root())
+        root_path = get_storage_root() if org_tree.partition == 'aws' else os.path.join(get_storage_root(), org_tree.partition)
+        _update_accounts_with_ou_path_map(org_tree.org_id, account_ou_map, root_path)
         logger.info('Updated currently stored Graphs with applicable AWS Organizations data')
 
         # create and cache a list of edges between all the accounts we have data for
         edge_list = []
         graph_objs = []
         for account in org_tree.accounts:
             try:
-                potential_path = os.path.join(get_storage_root(), account)
+                if org_tree.partition != 'aws':
+                    potential_path = get_default_graph_path(f'{org_tree.partition}:{account}')
+                else:
+                    potential_path = get_default_graph_path(account)
                 logger.debug('Trying to load a Graph from {}'.format(potential_path))
                 graph_obj = Graph.create_graph_from_local_disk(potential_path)
                 graph_objs.append(graph_obj)
@@ -175,12 +188,16 @@ def process_arguments(parsed_args: Namespace):
         org_tree.edge_list = edge_list
         logger.info('Compiled cross-account edges')
 
-        org_tree.save_organization_to_disk(os.path.join(get_storage_root(), org_tree.org_id))
+        if org_tree.partition != 'aws':
+            org_storage_path = get_default_graph_path(f'{org_tree.partition}:{org_tree.org_id}')
+        else:
+            org_storage_path = get_default_graph_path(org_tree.org_id)
+        org_tree.save_organization_to_disk(org_storage_path)
         logger.info('Stored organization data to disk')
 
     elif parsed_args.picked_orgs_cmd == 'display':
         # pull the existing data from disk
-        org_filepath = os.path.join(get_storage_root(), parsed_args.org)
+        org_filepath = get_default_graph_path(parsed_args.org)
         org_tree = OrganizationTree.create_from_dir(org_filepath)
 
         def _print_account(org_account: OrganizationAccount, indent_level: int, inherited_scps: List[Policy]):