reverse-diagrams 1.3.3__py3-none-any.whl → 1.3.5__py3-none-any.whl

src/aws/describe_organization.py

@@ -1,480 +1,359 @@
 """Describe Organizations."""
 import logging
 import os
+from typing import List, Dict, Any, Optional
+from pathlib import Path
 
 import emoji
 from colorama import Fore
 
+from .client_manager import get_client_manager
+from .exceptions import AWSServiceError
 from ..dgms.graph_mapper import create_file, create_mapper, find_ou_name
 from ..dgms.graph_template import graph_template
 from ..reports.save_results import save_results
-from .describe_sso import client
+from ..utils.progress import track_operation, get_progress_tracker
+from ..config import get_config
+from ..models import AWSAccount, OrganizationalUnit, AWSOrganization, validate_aws_response
 
-
-def describe_organization(region, org_client):
-    """
-    Describe the organization.
-
-    :param region: AWS Region
-    :return:
-    """
-    print(f"{Fore.GREEN}❇️ Describe Organization {Fore.RESET}")
-
-    organization = org_client.describe_organization()
-    organization = organization["Organization"]
-    return organization
-
-
-def list_roots(region):
-    """
-    List the roots.
-
-    :param region: AWS Region
-    :return:
-    """
-    org_client = client("organizations", region_name=region)
-    roots = org_client.list_roots()
-    return roots["Roots"]
-
-
-# def list organizational units pagination
-def list_organizational_units_pag(parent_id, region, next_token=None):
-    """
-    List organizational Units with pagination.
-
-    :param parent_id:
-    :param region:
-    :param next_token:
-    :return:
-    """
-    org_client = client("organizations", region_name=region)
-    paginator = org_client.get_paginator("list_organizational_units_for_parent")
-    response_iterator = paginator.paginate(
-        ParentId=parent_id,
-        PaginationConfig={"MaxItems": 1000, "PageSize": 4, "StartingToken": next_token},
-    )
-    response_iterator = response_iterator.build_full_result()
-    return response_iterator["OrganizationalUnits"]
-
-
-def list_organizational_units(parent_id, region, org_client, org_units=None):
-    """
-    List Organizational units.
-
-    :param org_client:
-    :param parent_id:
-    :param region:
-    :param org_units:
-    :return:
-    """
-
-    if org_units is None:
-        org_units = []
-    ous = org_client.list_organizational_units_for_parent(
-        ParentId=parent_id,
-        MaxResults=20,
-    )
-    ous = ous["OrganizationalUnits"]
-
-    if len(ous) >= 20:
-        logging.info("Paginating ...")
-        add_ous = list_organizational_units_pag(parent_id=parent_id,
-                                                region=region, next_token=ous["NextToken"]
-                                                )
-        for ou in add_ous:
-            ous.append(ou)
-        logging.debug(add_ous)
-
-    for o in ous:
-        org_units.append(o)
-    logging.debug(
-        f"The parent Id is: {parent_id}",
-    )
-    logging.debug(ous)
-    if len(ous) > 0:
-        for ou in ous:
-            logging.debug(ou)
-            if "Id" in ou.keys():
-                logging.debug(f"Search nested for: {ou['Name']}")
-                ous_next = list_organizational_units(parent_id=ou["Id"],
-                                                     region=region, org_units=org_units,
-                                                     org_client=org_client,
-                                                     )
-                logging.debug(ous_next)
-                if len(ous_next) > 0:
-                    logging.debug("Find Nested")
-
-    return org_units
-
-
-def list_parents(child_id, region):
-    """
-    List the parents of a child.
-
-    :param child_id:
-    :param region:
-    :return:
-    """
-    org_client = client("organizations", region_name=region)
-    response = org_client.list_parents(
-        ChildId=child_id,
-    )
-    return response["Parents"]
+logger = logging.getLogger(__name__)
 
 
-def index_ous(list_ous, region):
+def describe_organization(region: str) -> Dict[str, Any]:
     """
-    Index the parents of a child.
-
-    :param list_ous:
-    :param region:
-    :return list_ous:
+    Describe the organization.
 
-    """
-    org_client = client("organizations", region_name=region)
-    for ou in list_ous:
-        if "Id" in ou.keys() and len(ou) > 0:
-            response = org_client.list_parents(
-                ChildId=ou["Id"],
+    Args:
+        region: AWS Region
+
+    Returns:
+        Organization details
+
+    Raises:
+        AWSServiceError: If the API call fails
+    """
+    client_manager = get_client_manager(region=region)
+    progress = get_progress_tracker()
+    
+    progress.show_success("🏢 Getting Organization Info")
+    
+    try:
+        response = client_manager.call_api("organizations", "describe_organization")
+        validate_aws_response(response, ["Organization"])
+        
+        organization = response["Organization"]
+        logger.info(f"Retrieved organization: {organization.get('Id', 'Unknown')}")
+        return organization
+        
+    except Exception as e:
+        logger.error(f"Failed to describe organization: {e}")
+        raise AWSServiceError(f"Failed to describe organization: {e}")
+
+
+def list_roots(region: str) -> List[Dict[str, Any]]:
+    """
+    List the organization roots.
+
+    Args:
+        region: AWS Region
+
+    Returns:
+        List of root objects
+    """
+    client_manager = get_client_manager(region=region)
+    
+    try:
+        response = client_manager.call_api("organizations", "list_roots")
+        validate_aws_response(response, ["Roots"])
+        
+        roots = response["Roots"]
+        logger.info(f"Found {len(roots)} organization roots")
+        return roots
+        
+    except Exception as e:
+        logger.error(f"Failed to list roots: {e}")
+        raise AWSServiceError(f"Failed to list roots: {e}")
+
+
+def list_organizational_units(parent_id: str, region: str) -> List[Dict[str, Any]]:
+    """
+    List organizational units recursively.
+
+    Args:
+        parent_id: Parent ID to start from
+        region: AWS region
+
+    Returns:
+        List of all organizational units
+    """
+    client_manager = get_client_manager(region=region)
+    config = get_config()
+    all_ous = []
+    
+    def _get_ous_recursive(current_parent_id: str, depth: int = 0) -> None:
+        """Recursively get all OUs."""
+        if depth > 10:  # Prevent infinite recursion
+            logger.warning(f"Maximum recursion depth reached for parent {current_parent_id}")
+            return
+        
+        try:
+            # Get OUs for current parent
+            ous = client_manager.paginate_api_call(
+                "organizations",
+                "list_organizational_units_for_parent",
+                "OrganizationalUnits",
+                ParentId=current_parent_id,
+                PaginationConfig={
+                    'MaxItems': config.pagination.max_items,
+                    'PageSize': config.pagination.default_page_size
+                }
             )
-            logging.debug(response["Parents"])
-            ou["Parents"] = response["Parents"]
-    return list_ous
-
-
-def list_accounts_pag(region, next_token: str = None):
-    """
-    List accounts with pagination.
-
-    :param region:
-    :param next_token:
-    :return:
-    """
-    org_client = client("organizations", region_name=region)
-    paginator = org_client.get_paginator("list_accounts")
-    response_iterator = paginator.paginate(
-        PaginationConfig={"MaxItems": 1000, "PageSize": 20, "StartingToken": next_token}
-    )
-    response = response_iterator.build_full_result()
-    logging.info(response_iterator.build_full_result())
-    return response["Accounts"]
-
-
-def list_accounts(region, org_client):
-    """
-    List accounts.
-
-    :param org_client:
-    :param region:
-    :return:
-    """
-
-    accounts = org_client.list_accounts()
-    logging.info(accounts)
-    l_account = accounts["Accounts"]
-    logging.info(len(accounts["Accounts"]))
-    if len(accounts["Accounts"]) >= 20:
-        logging.info("Paginating ...")
-        ad_accounts = list_accounts_pag(region=region, next_token=accounts["NextToken"])
-        for ad in ad_accounts:
-            l_account.append(ad)
-        logging.info(f"You Organizations have {len(l_account)} Accounts")
-
-    return l_account
-
-
-def index_accounts(list_account, region, org_client):
-    """
-    Index accounts.
-
-    :param list_account:
-    :param region:
-    :return:
-    """
-    accounts = []
-
-    for a in list_account:
-        response = org_client.list_parents(
-            ChildId=a["Id"],
-        )
-
-        accounts.append(
-            {"account": a["Id"], "name": a["Name"], "parents": response["Parents"]}
-        )
-    return accounts
-
-
-# create organization complete map
-def find_ou_index(ous, search_id):
-    """
-    Find OU Name in list.
-
-    :param ous:
-    :param search_id:
-    :return:
-    """
-    for a in ous:
-        if a["Id"] == search_id:
-            return a
-
-
-# search ou in map
-def search_ou_map(map_ou: dict, ou_id, level=0, tree="."):
-    """
-    Search OU in map.
-
-    :param tree:
-    :param level:
-    :param map_ou:
-    :param ou_id:
-    :return:
-    """
-    for a in map_ou.keys():
-        # print(f'Searching {ou_id}... in {map_ou[a]["nestedOus"]}')
-
-        if len(map_ou[a]["nestedOus"]) > 0:
-            level += 1
-            tree += f".{a}"
-
-            if ou_id in map_ou[a]["nestedOus"].keys():
-                # search_ou_map(map_ou=map_ou[a]["nestedOus"], ou_id=ou_id, level=level, tree=tree)
-
-                return map_ou[a]
-            # else:
-            # search_ou_map(map_ou=map_ou[a]["nestedOus"], ou_id=ou_id, level=level, tree=tree)
-    return None
-
-
-def init_org_complete(
-        root_id,
-        org,
-        list_ous,
-):
-    """
-    Init organization dictionary.
-
-    :param root_id:
-    :param org:
-    :param list_ous:
-    :return:
+            
+            for ou in ous:
+                # Add parent information
+                parents_response = client_manager.call_api(
+                    "organizations",
+                    "list_parents",
+                    ChildId=ou["Id"]
+                )
+                ou["Parents"] = parents_response.get("Parents", [])
+                all_ous.append(ou)
+                
+                # Recursively get child OUs
+                _get_ous_recursive(ou["Id"], depth + 1)
+                
+        except Exception as e:
+            logger.warning(f"Failed to get OUs for parent {current_parent_id}: {e}")
+    
+    with track_operation("Listing organizational units") as task_id:
+        _get_ous_recursive(parent_id)
+    
+    logger.info(f"Found {len(all_ous)} organizational units")
+    return all_ous
+
+
+def list_accounts(region: str) -> List[Dict[str, Any]]:
+    """
+    List all accounts in the organization.
+
+    Args:
+        region: AWS region
+
+    Returns:
+        List of account objects with parent information
+    """
+    client_manager = get_client_manager(region=region)
+    config = get_config()
+    
+    with track_operation("Listing organization accounts") as task_id:
+        try:
+            # Get all accounts using pagination
+            accounts = client_manager.paginate_api_call(
+                "organizations",
+                "list_accounts",
+                "Accounts",
+                PaginationConfig={
+                    'MaxItems': config.pagination.max_items,
+                    'PageSize': config.pagination.default_page_size
+                }
+            )
+            
+            logger.info(f"Found {len(accounts)} accounts in organization")
+            
+            # Add parent information for each account
+            indexed_accounts = []
+            progress = get_progress_tracker()
+            
+            with progress.track_operation(f"Getting parent info for {len(accounts)} accounts", total=len(accounts)) as parent_task:
+                for i, account in enumerate(accounts):
+                    try:
+                        parents_response = client_manager.call_api(
+                            "organizations",
+                            "list_parents",
+                            ChildId=account["Id"]
+                        )
+                        
+                        indexed_accounts.append({
+                            "account": account["Id"],
+                            "name": account["Name"],
+                            "email": account.get("Email", ""),
+                            "status": account.get("Status", "ACTIVE"),
+                            "parents": parents_response.get("Parents", [])
+                        })
+                        
+                        progress.update_progress(parent_task)
+                        
+                    except Exception as e:
+                        logger.warning(f"Failed to get parents for account {account['Id']}: {e}")
+                        indexed_accounts.append({
+                            "account": account["Id"],
+                            "name": account["Name"],
+                            "email": account.get("Email", ""),
+                            "status": account.get("Status", "ACTIVE"),
+                            "parents": []
+                        })
+            
+            return indexed_accounts
+            
+        except Exception as e:
+            logger.error(f"Failed to list accounts: {e}")
+            raise AWSServiceError(f"Failed to list accounts: {e}")
+
+
+def create_organization_complete_map(
+    root_id: str,
+    organization: Dict[str, Any],
+    organizational_units: List[Dict[str, Any]],
+    accounts: List[Dict[str, Any]]
+) -> Dict[str, Any]:
+    """
+    Create a complete organization map with nested structure.
+
+    Args:
+        root_id: Organization root ID
+        organization: Organization details
+        organizational_units: List of OUs
+        accounts: List of accounts
+
+    Returns:
+        Complete organization structure
     """
     organizations_complete = {
         "rootId": root_id,
-        "masterAccountId": org["MasterAccountId"],
+        "masterAccountId": organization["MasterAccountId"],
         "noOutAccounts": [],
         "organizationalUnits": {},
     }
-    # Iterate in ous for getting ous tree
-    for a, i in zip(list_ous, range(len(list_ous))):
-        for p in a["Parents"]:
-            if p["Type"] == "ROOT":
-                organizations_complete["organizationalUnits"][a["Name"]] = {
-                    "Id": a["Id"],
-                    "Name": a["Name"],
+    
+    # Build OU hierarchy
+    for ou in organizational_units:
+        for parent in ou.get("Parents", []):
+            if parent["Type"] == "ROOT":
+                organizations_complete["organizationalUnits"][ou["Name"]] = {
+                    "Id": ou["Id"],
+                    "Name": ou["Name"],
                     "accounts": {},
                     "nestedOus": {},
                 }
-    return organizations_complete
-
-
-# create organization complete map
-def map_organizations_complete(
-        organizations_complete: dict,
-        list_ous,
-        llist_accounts,
-        reference_outs_list,
-):
-    """
-    Create complete mapper file.
-
-    :param reference_outs_list:
-    :param organizations_complete:
-    :param list_ous:
-    :param llist_accounts:
-    :return:
-    """
-    # Iterate in ous for getting ous tree
-    for a, i in zip(list_ous, range(len(list_ous))):
-        for p in a["Parents"]:
-            if p["Type"] == "ORGANIZATIONAL_UNIT":
-                o = find_ou_name(reference_outs_list, p["Id"])
-
-                if o not in organizations_complete["organizationalUnits"].keys():
-                    p = search_ou_map(
-                        organizations_complete["organizationalUnits"], ou_id=o
-                    )
-                    o = p["Name"]
-
-                organizations_complete["organizationalUnits"][o]["nestedOus"][
-                    find_ou_name(reference_outs_list, a["Id"])
-                ] = {"Id": a["Id"], "Name": a["Name"], "accounts": [], "nestedOus": {}}
-                # print(organizations_complete["organizationalUnits"][o]["nestedOus"])
-                if (
-                        len(organizations_complete["organizationalUnits"][o]["nestedOus"])
-                        > 0
-                ):
-                    new_list_ous = organizations_complete["organizationalUnits"][o][
-                        "nestedOus"
-                    ]
-
-                    new_list_ous = plop_dict_out(ous_list=list_ous, ou=new_list_ous)
-                    organizations_complete = map_organizations_complete(
-                        organizations_complete=organizations_complete,
-                        list_ous=new_list_ous,
-                        llist_accounts=llist_accounts,
-                        reference_outs_list=reference_outs_list,
-                    )
-
-    return organizations_complete
-
-
-def plop_dict_out(
-        ous_list: list,
-        ou,
-):
-    """
-    Clean list.
-
-    :param ous_list:
-    :param ou:
-    :return:
-    """
-    for o in ou.keys():
-        # for c in ou.keys():
-        for unit in ous_list:
-            if unit["Id"] == ou[o]["Id"]:
-                ous_list.remove(unit)
-
-    return ous_list
-
-
-def set_accounts_tree(llist_accounts, organizations_complete, list_ous):
-    """
-    Set accounts tree.
-
-    :param llist_accounts:
-    :param organizations_complete:
-    :param list_ous:
-    :return:
-    """
-    # Iterate in list accounts to get parent ous
-    for c, i in zip(llist_accounts, range(len(llist_accounts))):
-        # print(f"\n    aa_{i}= OrganizationsAccount(\"{c['account']}\")", file=f)
-        for p in c["parents"]:
-            if p["Type"] == "ROOT":
-                organizations_complete["noOutAccounts"].append(
-                    {"account": c["account"], "name": c["name"]}
-                )
-
-            for o, j in zip(list_ous, range(len(list_ous))):
-                if p["Id"] == o["Id"] and p["Type"] == "ORGANIZATIONAL_UNIT":
-                    organizations_complete["organizationalUnits"][
-                        find_ou_name(list_ous, o["Id"])
-                    ]["accounts"][c["name"]] = {
-                        "account": c["account"],
-                        "name": c["name"],
+    
+    # Add accounts to appropriate OUs or root
+    for account in accounts:
+        for parent in account.get("parents", []):
+            if parent["Type"] == "ROOT":
+                organizations_complete["noOutAccounts"].append({
+                    "account": account["account"],
+                    "name": account["name"]
+                })
+            elif parent["Type"] == "ORGANIZATIONAL_UNIT":
+                # Find the OU name
+                ou_name = find_ou_name(organizational_units, parent["Id"])
+                if ou_name and ou_name in organizations_complete["organizationalUnits"]:
+                    organizations_complete["organizationalUnits"][ou_name]["accounts"][account["name"]] = {
+                        "account": account["account"],
+                        "name": account["name"]
                     }
-
+    
     return organizations_complete
 
 
-def graph_organizations(diagrams_path, region, auto):
+def graph_organizations(diagrams_path: str, region: str, auto: bool) -> None:
     """
-    Create organizations Graph.
+    Create organizations graph with improved error handling and progress tracking.
 
-    :param auto:
-    :param diagrams_path:
-    :param region:
-    :return:
+    Args:
+        diagrams_path: Output directory path
+        region: AWS region
+        auto: Whether to automatically create diagrams
     """
     template_file = "graph_org.py"
-    code_path = f"{diagrams_path}/code"
-    json_path = f"{diagrams_path}/json"
-    create_file(
-        template_content=graph_template,
-        file_name=template_file,
-        directory_path=code_path,
-    )
-    org_client = client("organizations", region_name=region)
-    organization = describe_organization(region=region, org_client=org_client)
-    print(Fore.BLUE + emoji.emojize(":sparkle: Getting Organization Info" + Fore.RESET))
-    logging.debug(organization)
-    logging.debug("The Roots Info")
-    roots = list_roots(region=region)
-    logging.debug(roots)
-
-    print(
-        Fore.BLUE + emoji.emojize(":sparkle: List Organizational Units " + Fore.RESET)
-    )
-    logging.debug("The Organizational Units list ")
-
-    ous = list_organizational_units(parent_id=roots[0]["Id"], region=region, org_client=org_client)
-    logging.debug(ous)
-    logging.debug("The Organizational Units list with parents info")
-    i_ous = index_ous(ous, region=region)
-    logging.debug(i_ous)
-
-    print(
-        Fore.BLUE
-        + emoji.emojize(":sparkle: Getting the Account list info" + Fore.RESET)
-    )
-
-    l_accounts = list_accounts(region=region, org_client=org_client)
-    logging.debug(l_accounts)
-    logging.debug("The Account list with parents info")
-
-    print(
-        Fore.YELLOW
-        + emoji.emojize(
-            f":information:  There are {len(l_accounts)} Accounts in your organization"
-            + Fore.RESET
+    code_path = Path(diagrams_path) / "code"
+    json_path = Path(diagrams_path) / "json"
+    
+    # Ensure directories exist
+    code_path.mkdir(parents=True, exist_ok=True)
+    json_path.mkdir(parents=True, exist_ok=True)
+    
+    progress = get_progress_tracker()
+    
+    try:
+        # Create template file
+        create_file(
+            template_content=graph_template,
+            file_name=template_file,
+            directory_path=str(code_path),
         )
-    )
-
-    i_accounts = index_accounts(l_accounts, region=region, org_client=org_client)
-    logging.debug(i_accounts)
-
-    file_name = "organizations.json"
-    save_results(results=i_accounts, filename=file_name, directory_path=json_path)
-
-    create_mapper(
-        template_file=f"{code_path}/{template_file}",
-        org=organization,
-        root_id=roots[0]["Id"],
-        list_ous=ous,
-        list_accounts=i_accounts,
-    )
-    file_name = "organizations_complete.json"
-    # view in console
-    organizations_complete_f = map_organizations_complete(
-        organizations_complete=init_org_complete(
-            org=organization, root_id=roots[0]["Id"], list_ous=i_ous
-        ),
-        llist_accounts=i_accounts,
-        list_ous=i_ous,
-        reference_outs_list=i_ous.copy(),
-    )
-    organizations_complete_f = (
-        set_accounts_tree(
-            llist_accounts=i_accounts,
-            organizations_complete=organizations_complete_f,
-            list_ous=i_ous,
-        ),
-    )
-
-    save_results(results=organizations_complete_f, filename=f"{json_path}/{file_name}")
-
-    if auto:
-        print(f"{Fore.GREEN}❇️ Creating diagrams in {code_path}")
-        command = os.system(f"cd {code_path} && python3 {template_file}")
-        if command != 0:
-            print(Fore.RED + "❌ Error creating diagrams")
-            print(command)
-
-    else:
-        print(
-            Fore.YELLOW
-            + emoji.emojize(
-                f":sparkles: Run -> python3 {code_path}/graph_org.py " + Fore.RESET
-            )
+        
+        # Get organization data
+        organization = describe_organization(region=region)
+        roots = list_roots(region=region)
+        
+        if not roots:
+            raise AWSServiceError("No organization roots found")
+        
+        root_id = roots[0]["Id"]
+        
+        # Get organizational units
+        progress.show_success("📋 Listing Organizational Units")
+        organizational_units = list_organizational_units(parent_id=root_id, region=region)
+        
+        # Get accounts
+        progress.show_success("👥 Getting Account Information")
+        accounts = list_accounts(region=region)
+        
+        progress.show_summary(
+            "Organization Summary",
+            [
+                f"Organization ID: {organization.get('Id', 'Unknown')}",
+                f"Master Account: {organization.get('MasterAccountId', 'Unknown')}",
+                f"Organizational Units: {len(organizational_units)}",
+                f"Accounts: {len(accounts)}"
+            ]
+        )
+        
+        # Save basic accounts data
+        save_results(
+            results=accounts,
+            filename="organizations.json",
+            directory_path=str(json_path)
+        )
+        
+        # Create diagram mapper
+        create_mapper(
+            template_file=str(code_path / template_file),
+            org=organization,
+            root_id=root_id,
+            list_ous=organizational_units,
+            list_accounts=accounts,
         )
+        
+        # Create complete organization map
+        organizations_complete = create_organization_complete_map(
+            root_id=root_id,
+            organization=organization,
+            organizational_units=organizational_units,
+            accounts=accounts
+        )
+        
+        save_results(
+            results=organizations_complete,
+            filename="organizations_complete.json",
+            directory_path=str(json_path)
+        )
+        
+        if auto:
+            progress.show_success(f"🎨 Creating diagrams in {code_path}")
+            command = os.system(f"cd {code_path} && python3 {template_file}")
+            if command != 0:
+                progress.show_error("Failed to create diagrams", f"Command exit code: {command}")
+            else:
+                progress.show_success("Diagrams created successfully")
+        else:
+            progress.show_success(
+                "Ready to create diagrams",
+                f"Run: python3 {code_path / template_file}"
+            )
+            
+    except Exception as e:
+        logger.error(f"Failed to create organization graph: {e}")
+        progress.show_error("Organization diagram generation failed", str(e))
+        raise