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

src/plugins/builtin/identity_center_plugin.py

@@ -0,0 +1,496 @@
+"""AWS IAM Identity Center (SSO) plugin for generating identity and access diagrams."""
+import logging
+from typing import Dict, Any, List
+from pathlib import Path
+
+from src.plugins.base import AWSServicePlugin, PluginMetadata
+from src.aws.client_manager import AWSClientManager
+from src.models import DiagramConfig
+from src.utils.concurrent import get_concurrent_processor
+from src.utils.progress import get_progress_tracker
+
+logger = logging.getLogger(__name__)
+
+
+class IdentityCenterPlugin(AWSServicePlugin):
+    """Plugin for AWS IAM Identity Center (SSO) diagram generation."""
+    
+    @property
+    def metadata(self) -> PluginMetadata:
+        """Get plugin metadata."""
+        return PluginMetadata(
+            name="identity-center",
+            version="1.0.0",
+            description="Generate diagrams for AWS IAM Identity Center (SSO) groups, users, and permissions",
+            author="Reverse Diagrams Team",
+            aws_services=["sso-admin", "identitystore", "organizations"],
+            dependencies=[]
+        )
+    
+    def collect_data(self, client_manager: AWSClientManager, region: str, **kwargs) -> Dict[str, Any]:
+        """
+        Collect AWS IAM Identity Center data.
+        
+        Args:
+            client_manager: AWS client manager
+            region: AWS region
+            **kwargs: Additional parameters
+            
+        Returns:
+            Dictionary containing Identity Center data
+        """
+        logger.debug(f"Collecting AWS IAM Identity Center data from region {region}")
+        progress = get_progress_tracker()
+        
+        data = {
+            "region": region,
+            "sso_instances": [],
+            "identity_store_id": "",
+            "instance_arn": "",
+            "groups": [],
+            "users": [],
+            "group_memberships": [],
+            "permission_sets": [],
+            "permission_set_details": {},
+            "account_assignments": [],
+            "accounts": [],
+            "final_account_assignments": {}
+        }
+        
+        try:
+            # Get SSO instances
+            progress.show_success("🔐 Getting Identity Store Instance Info")
+            instances_response = client_manager.call_api("sso-admin", "list_instances")
+            data["sso_instances"] = instances_response.get("Instances", [])
+            
+            if not data["sso_instances"]:
+                raise ValueError("No SSO instances found")
+            
+            data["identity_store_id"] = data["sso_instances"][0]["IdentityStoreId"]
+            data["instance_arn"] = data["sso_instances"][0]["InstanceArn"]
+            
+            logger.debug(f"Using Identity Store ID: {data['identity_store_id']}")
+            
+            # Get groups
+            progress.show_success("👥 Listing Groups")
+            data["groups"] = self._list_groups(client_manager, data["identity_store_id"])
+            logger.debug(f"Found {len(data['groups'])} groups")
+            
+            # Get users
+            progress.show_success("👤 Listing Users")
+            data["users"] = self._list_users(client_manager, data["identity_store_id"])
+            logger.debug(f"Found {len(data['users'])} users")
+            
+            # Get group memberships
+            progress.show_success("🔗 Getting Group Memberships")
+            data["group_memberships"] = self._get_group_memberships(
+                client_manager, data["identity_store_id"], data["groups"]
+            )
+            
+            # Complete group members with user information
+            data["group_memberships"] = self._complete_group_members(
+                data["group_memberships"], data["users"]
+            )
+            
+            # Get permission sets
+            progress.show_success("🛡️ Listing Permission Sets")
+            data["permission_sets"] = self._list_permission_sets(
+                client_manager, data["instance_arn"]
+            )
+            logger.debug(f"Found {len(data['permission_sets'])} permission sets")
+            
+            # Get permission set details
+            data["permission_set_details"] = self._get_permission_set_details(
+                client_manager, data["instance_arn"], data["permission_sets"]
+            )
+            
+            # Get organization accounts
+            progress.show_success("🏢 Getting Organization Accounts")
+            data["accounts"] = self._list_organization_accounts(client_manager)
+            logger.debug(f"Found {len(data['accounts'])} accounts")
+            
+            # Get account assignments
+            progress.show_success("📋 Getting Account Assignments")
+            data["account_assignments"] = self._get_account_assignments(
+                client_manager,
+                data["instance_arn"],
+                data["accounts"],
+                data["permission_sets"]
+            )
+            
+            # Add user and group information to assignments
+            data["account_assignments"] = self._enrich_account_assignments(
+                data["account_assignments"],
+                data["group_memberships"],
+                data["users"],
+                data["permission_set_details"]
+            )
+            
+            # Create final account assignments structure
+            data["final_account_assignments"] = self._organize_account_assignments(
+                data["accounts"], data["account_assignments"]
+            )
+            
+            progress.show_summary(
+                "Identity Center Summary",
+                [
+                    f"Identity Store ID: {data['identity_store_id']}",
+                    f"Groups: {len(data['groups'])}",
+                    f"Users: {len(data['users'])}",
+                    f"Permission Sets: {len(data['permission_sets'])}",
+                    f"Account Assignments: {len(data['account_assignments'])}"
+                ]
+            )
+            
+        except Exception as e:
+            logger.error(f"Failed to collect Identity Center data: {e}")
+            raise
+        
+        return data
+    
+    def _list_groups(self, client_manager: AWSClientManager, identity_store_id: str) -> List[Dict[str, Any]]:
+        """List all groups in the identity store."""
+        try:
+            groups = client_manager.paginate_api_call(
+                "identitystore",
+                "list_groups",
+                "Groups",
+                IdentityStoreId=identity_store_id
+            )
+            return groups
+        except Exception as e:
+            logger.error(f"Failed to list groups: {e}")
+            return []
+    
+    def _list_users(self, client_manager: AWSClientManager, identity_store_id: str) -> List[Dict[str, Any]]:
+        """List all users in the identity store."""
+        try:
+            users = client_manager.paginate_api_call(
+                "identitystore",
+                "list_users",
+                "Users",
+                IdentityStoreId=identity_store_id
+            )
+            return users
+        except Exception as e:
+            logger.error(f"Failed to list users: {e}")
+            return []
+    
+    def _get_group_memberships(
+        self, 
+        client_manager: AWSClientManager, 
+        identity_store_id: str, 
+        groups: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Get group memberships for all groups."""
+        group_memberships = []
+        progress = get_progress_tracker()
+        
+        with progress.track_operation(f"Getting memberships for {len(groups)} groups", total=len(groups)) as task_id:
+            for group in groups:
+                try:
+                    memberships = client_manager.paginate_api_call(
+                        "identitystore",
+                        "list_group_memberships",
+                        "GroupMemberships",
+                        IdentityStoreId=identity_store_id,
+                        GroupId=group["GroupId"]
+                    )
+                    
+                    group_memberships.append({
+                        "group_id": group["GroupId"],
+                        "group_name": group["DisplayName"],
+                        "members": memberships
+                    })
+                    
+                    progress.update_progress(task_id)
+                    
+                except Exception as e:
+                    logger.warning(f"Failed to get memberships for group {group['GroupId']}: {e}")
+                    group_memberships.append({
+                        "group_id": group["GroupId"],
+                        "group_name": group["DisplayName"],
+                        "members": []
+                    })
+        
+        return group_memberships
+    
+    def _complete_group_members(
+        self, 
+        group_memberships: List[Dict[str, Any]], 
+        users: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Complete group member information with user details."""
+        user_lookup = {user["UserId"]: user for user in users}
+        
+        for group_membership in group_memberships:
+            for member in group_membership["members"]:
+                user_id = member.get("MemberId", {}).get("UserId")
+                if user_id and user_id in user_lookup:
+                    user = user_lookup[user_id]
+                    member["MemberId"]["UserName"] = user.get("UserName", "Unknown")
+        
+        return group_memberships
+    
+    def _list_permission_sets(self, client_manager: AWSClientManager, instance_arn: str) -> List[str]:
+        """List all permission sets."""
+        try:
+            permission_sets = client_manager.paginate_api_call(
+                "sso-admin",
+                "list_permission_sets",
+                "PermissionSets",
+                InstanceArn=instance_arn
+            )
+            return permission_sets
+        except Exception as e:
+            logger.error(f"Failed to list permission sets: {e}")
+            return []
+    
+    def _get_permission_set_details(
+        self, 
+        client_manager: AWSClientManager, 
+        instance_arn: str, 
+        permission_sets: List[str]
+    ) -> Dict[str, str]:
+        """Get detailed information for permission sets."""
+        permission_set_details = {}
+        progress = get_progress_tracker()
+        
+        with progress.track_operation(f"Getting details for {len(permission_sets)} permission sets", total=len(permission_sets)) as task_id:
+            for permission_set_arn in permission_sets:
+                try:
+                    response = client_manager.call_api(
+                        "sso-admin",
+                        "describe_permission_set",
+                        InstanceArn=instance_arn,
+                        PermissionSetArn=permission_set_arn
+                    )
+                    
+                    permission_set = response.get("PermissionSet", {})
+                    permission_set_details[permission_set_arn] = permission_set.get("Name", "Unknown")
+                    
+                    progress.update_progress(task_id)
+                    
+                except Exception as e:
+                    logger.warning(f"Failed to describe permission set {permission_set_arn}: {e}")
+                    permission_set_details[permission_set_arn] = "Unknown"
+        
+        return permission_set_details
+    
+    def _list_organization_accounts(self, client_manager: AWSClientManager) -> List[Dict[str, Any]]:
+        """List organization accounts."""
+        try:
+            accounts = client_manager.paginate_api_call(
+                "organizations",
+                "list_accounts",
+                "Accounts"
+            )
+            return accounts
+        except Exception as e:
+            logger.warning(f"Failed to list organization accounts: {e}")
+            return []
+    
+    def _get_account_assignments(
+        self,
+        client_manager: AWSClientManager,
+        instance_arn: str,
+        accounts: List[Dict[str, Any]],
+        permission_sets: List[str]
+    ) -> List[Dict[str, Any]]:
+        """Get account assignments for all accounts and permission sets."""
+        all_assignments = []
+        progress = get_progress_tracker()
+        
+        total_operations = len(accounts) * len(permission_sets)
+        
+        with progress.track_operation(f"Getting account assignments", total=total_operations) as task_id:
+            for account in accounts:
+                for permission_set_arn in permission_sets:
+                    try:
+                        assignments = client_manager.paginate_api_call(
+                            "sso-admin",
+                            "list_account_assignments",
+                            "AccountAssignments",
+                            InstanceArn=instance_arn,
+                            AccountId=account["Id"],
+                            PermissionSetArn=permission_set_arn
+                        )
+                        
+                        all_assignments.extend(assignments)
+                        progress.update_progress(task_id)
+                        
+                    except Exception as e:
+                        logger.debug(f"No assignments for account {account['Id']} and permission set {permission_set_arn}: {e}")
+                        progress.update_progress(task_id)
+        
+        return all_assignments
+    
+    def _enrich_account_assignments(
+        self,
+        account_assignments: List[Dict[str, Any]],
+        group_memberships: List[Dict[str, Any]],
+        users: List[Dict[str, Any]],
+        permission_set_details: Dict[str, str]
+    ) -> List[Dict[str, Any]]:
+        """Add user and group information to account assignments."""
+        # Create lookup dictionaries
+        group_lookup = {group["group_id"]: group for group in group_memberships}
+        user_lookup = {user["UserId"]: user for user in users}
+        
+        progress = get_progress_tracker()
+        
+        with progress.track_operation(f"Enriching {len(account_assignments)} assignments", total=len(account_assignments)) as task_id:
+            for assignment in account_assignments:
+                # Add permission set name
+                permission_set_arn = assignment.get("PermissionSetArn", "")
+                assignment["PermissionSetName"] = permission_set_details.get(permission_set_arn, "Unknown")
+                
+                # Add group or user information
+                principal_type = assignment.get("PrincipalType", "")
+                principal_id = assignment.get("PrincipalId", "")
+                
+                if principal_type == "GROUP" and principal_id in group_lookup:
+                    assignment["GroupName"] = group_lookup[principal_id]["group_name"]
+                elif principal_type == "USER" and principal_id in user_lookup:
+                    assignment["UserName"] = user_lookup[principal_id].get("UserName", "Unknown")
+                
+                progress.update_progress(task_id)
+        
+        return account_assignments
+    
+    def _organize_account_assignments(
+        self, 
+        accounts: List[Dict[str, Any]], 
+        account_assignments: List[Dict[str, Any]]
+    ) -> Dict[str, List[Dict[str, Any]]]:
+        """Organize account assignments by account name."""
+        account_lookup = {account["Id"]: account["Name"] for account in accounts}
+        organized_assignments = {}
+        
+        for assignment in account_assignments:
+            account_id = assignment.get("AccountId", "")
+            account_name = account_lookup.get(account_id, account_id)
+            
+            if account_name not in organized_assignments:
+                organized_assignments[account_name] = []
+            
+            organized_assignments[account_name].append(assignment)
+        
+        return organized_assignments
+    
+    def generate_diagram_code(self, data: Dict[str, Any], config: DiagramConfig) -> str:
+        """
+        Generate diagram code for AWS IAM Identity Center.
+        
+        Args:
+            data: Identity Center data collected from AWS
+            config: Diagram configuration
+            
+        Returns:
+            Python code for generating Identity Center diagram
+        """
+        logger.debug("Generating AWS IAM Identity Center diagram code")
+        
+        code_lines = [
+            "from diagrams import Diagram, Cluster, Edge",
+            "from diagrams.aws.management import Organizations, OrganizationsAccount, OrganizationsOrganizationalUnit",
+            "from diagrams.aws.general import Users, User",
+            "from diagrams.aws.security import IAMPermissions",
+            "",
+            f'with Diagram("{config.title}", show=False, direction="{config.direction}"):'
+        ]
+        
+        # Get data
+        final_assignments = data.get("final_account_assignments", {})
+        group_memberships = data.get("group_memberships", [])
+        
+        # Create group lookup for members
+        group_lookup = {group["group_name"]: group for group in group_memberships}
+        
+        # Generate account assignment clusters
+        for account_name, assignments in final_assignments.items():
+            if not assignments:
+                continue
+                
+            code_lines.append(f"    with Cluster('Account: {account_name}'):")
+            
+            # Group assignments by group/user
+            processed_principals = set()
+            
+            for assignment in assignments:
+                principal_key = None
+                
+                if "GroupName" in assignment and assignment["GroupName"] not in processed_principals:
+                    group_name = assignment["GroupName"]
+                    principal_key = f"group_{self._format_name_for_code(group_name)}"
+                    processed_principals.add(group_name)
+                    
+                    code_lines.append(f"        with Cluster('Group: {group_name}'):")
+                    code_lines.append(f"            {principal_key} = Users(\"{self._split_long_name(group_name)}\")")
+                    code_lines.append(f"            {principal_key} \\")
+                    code_lines.append(f"                - Edge(color=\"brown\", style=\"dotted\", label=\"Permissions Set\") \\")
+                    code_lines.append(f"                - IAMPermissions(\"{self._split_long_name(assignment['PermissionSetName'])}\")")
+                    
+                    # Add group members if available
+                    if group_name in group_lookup:
+                        members = group_lookup[group_name]["members"]
+                        if members:
+                            member_names = [m.get("MemberId", {}).get("UserName", "Unknown") for m in members]
+                            members_code = self._create_users_list(member_names)
+                            code_lines.append(f"            members_{self._format_name_for_code(group_name)} = {members_code}")
+                            code_lines.append(f"            {principal_key} \\")
+                            code_lines.append(f"                - Edge(color=\"darkgreen\", style=\"dotted\", label=\"Member\") \\")
+                            code_lines.append(f"                - members_{self._format_name_for_code(group_name)}")
+                
+                elif "UserName" in assignment and assignment["UserName"] not in processed_principals:
+                    user_name = assignment["UserName"]
+                    principal_key = f"user_{self._format_name_for_code(user_name)}"
+                    processed_principals.add(user_name)
+                    
+                    code_lines.append(f"        with Cluster('User: {user_name}'):")
+                    code_lines.append(f"            {principal_key} = User(\"{self._split_long_name(user_name)}\")")
+                    code_lines.append(f"            {principal_key} \\")
+                    code_lines.append(f"                - Edge(color=\"brown\", style=\"dotted\") \\")
+                    code_lines.append(f"                - IAMPermissions(\"{self._split_long_name(assignment['PermissionSetName'])}\")")
+        
+        return "\n".join(code_lines)
+    
+    def _format_name_for_code(self, name: str) -> str:
+        """Format name for use in Python code."""
+        import re
+        # Remove special characters and spaces
+        formatted = re.sub(r'[^\w]', '', name)
+        return formatted if formatted else "Unknown"
+    
+    def _split_long_name(self, name: str) -> str:
+        """Split long names for better display."""
+        if len(name) > 17:
+            return name[:16] + "\\n" + name[16:]
+        return name
+    
+    def _create_users_list(self, user_names: List[str]) -> str:
+        """Create a list of User objects for diagram code."""
+        if not user_names:
+            return "[]"
+        
+        user_objects = []
+        for user_name in user_names[:5]:  # Limit to 5 users to avoid clutter
+            user_objects.append(f'User("{self._split_long_name(user_name)}")')
+        
+        if len(user_names) > 5:
+            user_objects.append(f'User("... and {len(user_names) - 5} more")')
+        
+        return "[" + ", ".join(user_objects) + "]"
+    
+    def get_required_permissions(self) -> List[str]:
+        """Get required AWS permissions for Identity Center plugin."""
+        return [
+            "sso:ListInstances",
+            "sso:ListPermissionSets",
+            "sso:DescribePermissionSet",
+            "sso:ListAccountAssignments",
+            "identitystore:ListGroups",
+            "identitystore:ListUsers",
+            "identitystore:ListGroupMemberships",
+            "organizations:ListAccounts"
+        ]