fix privileges issue

Author: tmunzer-AIDE

src/mistapi/__api_session.py

@@ -948,11 +948,7 @@ def get_privilege_by_org_id(self, org_id: str):
         """
         logger.debug("apisession:get_privilege_by_org_id")
         org_priv = next(
-            (
-                priv
-                for priv in self.privileges.privileges
-                if priv.get("org_id") == org_id
-            ),
+            (priv for priv in self.privileges if priv.get("org_id") == org_id),
             None,
         )
         if org_priv:
@@ -1001,7 +997,7 @@ def get_privilege_by_org_id(self, org_id: str):
                 msp_priv = next(
                     (
                         priv
-                        for priv in self.privileges.privileges
+                        for priv in self.privileges
                         if priv.get("scope") == "msp" and priv.get("msp_id") == msp_id
                     ),
                     None,

src/mistapi/__models/privilege.py

@@ -1,41 +1,8 @@
-from typing import Any
+from typing import Any, Iterator
 
 from tabulate import tabulate
 
 
-class Privileges:
-    def __init__(self, privileges: list[dict]) -> None:
-        self.privileges: list[_Privilege] = []
-        for privilege in privileges:
-            self.privileges.append(_Privilege(privilege))
-
-    def __str__(self) -> str:
-        columns_headers = [
-            "scope",
-            "role",
-            "name",
-            "site_id",
-            "org_name",
-            "org_id",
-            "msp_name",
-            "msp_id",
-            "views",
-        ]
-        table = []
-        for entry in self.privileges:
-            temp = []
-            for field in columns_headers:
-                if hasattr(entry, field):
-                    temp.append(str(getattr(entry, field)))
-                else:
-                    temp.append("")
-            table.append(temp)
-        return tabulate(table, columns_headers)
-
-    def display(self):
-        return str(self)
-
-
 class _Privilege:
     def __init__(self, privilege) -> None:
         self.scope: str = ""
@@ -77,3 +44,42 @@ def get(self, key: str, default: Any | None = None) -> Any:
             return getattr(self, key)
         else:
             return default
+
+
+class Privileges:
+    def __init__(self, privileges: list[dict]) -> None:
+        self.privileges: list[_Privilege] = []
+        for privilege in privileges:
+            self.privileges.append(_Privilege(privilege))
+
+    def __iter__(self) -> Iterator[_Privilege]:
+        """Return an iterator over the privileges."""
+        if not self.privileges:
+            return iter([])
+        return iter(self.privileges)
+
+    def __str__(self) -> str:
+        columns_headers = [
+            "scope",
+            "role",
+            "name",
+            "site_id",
+            "org_name",
+            "org_id",
+            "msp_name",
+            "msp_id",
+            "views",
+        ]
+        table = []
+        for entry in self.privileges:
+            temp = []
+            for field in columns_headers:
+                if hasattr(entry, field):
+                    temp.append(str(getattr(entry, field)))
+                else:
+                    temp.append("")
+            table.append(temp)
+        return tabulate(table, columns_headers)
+
+    def display(self):
+        return str(self)

src/mistapi/cli.py

@@ -20,6 +20,7 @@
 import mistapi
 from mistapi.__api_response import APIResponse as _APIResponse
 from mistapi.__logger import console
+from mistapi.__models.privilege import Privileges
 
 
 ###########################################
@@ -63,7 +64,7 @@ def wrapper(*args, **kwargs):
 
 ###########################################
 #### CLI SELECTIONS
-def _forge_privileges(mist_session: mistapi.APISession, msp_id: str):
+def _forge_privileges(mist_session: mistapi.APISession, msp_id: str) -> Privileges:
     """
     Function to generate user privileges for Orgs belonging to a MSP Account
 
@@ -84,11 +85,11 @@ def _forge_privileges(mist_session: mistapi.APISession, msp_id: str):
     custom_privileges = []
     for org in orgs:
         custom_privileges.append(mist_session.get_privilege_by_org_id(org["id"]))
-    return custom_privileges
+    return Privileges(custom_privileges)
 
 
 @is_authenticated
-def _select_msp(mist_session: mistapi.APISession) -> list:
+def _select_msp(mist_session: mistapi.APISession) -> Privileges:
     """
     Function to list all the Mist MSPs allowed for the current user
     and ask to pick one. Return the list org ORG privileges based
@@ -105,12 +106,10 @@ def _select_msp(mist_session: mistapi.APISession) -> list:
         List of ORG privileges
     """
     msp_accounts = [
-        priv
-        for priv in mist_session.privileges.privileges
-        if priv.get("scope") == "msp"
+        priv for priv in mist_session.privileges if priv.get("scope") == "msp"
     ]
     if len(msp_accounts) == 0:
-        return mist_session.privileges.privileges
+        return mist_session.privileges
     else:
         msp_accounts = sorted(msp_accounts, key=lambda x: x.get("name").lower())
         while True:
@@ -128,16 +127,16 @@ def _select_msp(mist_session: mistapi.APISession) -> list:
             if resp == "q":
                 sys.exit(0)
             elif resp.lower() == "n":
-                standalone = []
-                for priv in mist_session.privileges.privileges:
+                standalone: list = []
+                for priv in mist_session.privileges:
                     msp = [
                         msp
                         for msp in msp_accounts
                         if msp.get("msp_id") == priv.get("msp_id", "xyz")
                     ]
                     if not msp:
                         standalone.append(priv)
-                return standalone
+                return Privileges(standalone)
                 # return [priv for priv in mist_session.privileges if not priv.get("msp_id")]
             else:
                 tested_val = _test_choice(resp, i)
@@ -172,41 +171,52 @@ def select_org(mist_session: mistapi.APISession, allow_many=False) -> list:
     """
     data = _select_msp(mist_session)
     data = [d for d in data if d.get("name")]
-    data = sorted(data, key=lambda x: x["name"].lower())
+    data = sorted(data, key=lambda x: x.get("name").lower())
     while True:
         i = -1
         org_ids: list[str] = []
         resp_ids: list[str] = []
         print("\r\nAvailable organizations:")
         for privilege in data:
-            if privilege["scope"] == "org" and privilege["org_id"] not in org_ids:
+            if (
+                privilege.get("scope") == "org"
+                and privilege.get("org_id") not in org_ids
+            ):
                 i += 1
-                org_ids.append(privilege["org_id"])
-                print(f"{i}) {privilege['name']} (id: {privilege['org_id']})")
+                org_ids.append(privilege.get("org_id"))
+                print(f"{i}) {privilege.get('name')} (id: {privilege.get('org_id')})")
 
         orgs_with_sites: list[dict] = []
         for privilege in data:
-            if privilege["scope"] == "site" and privilege["org_id"] not in org_ids:
-                index = _search_org(orgs_with_sites, privilege["org_id"])
+            if (
+                privilege.get("scope") == "site"
+                and privilege.get("org_id") not in org_ids
+            ):
+                index = _search_org(orgs_with_sites, privilege.get("org_id"))
                 if index is None:
                     i += 1
-                    org_ids.append(privilege["org_id"])
-                    print(f"{i}) {privilege['name']} (id: {privilege['org_id']})")
+                    org_ids.append(privilege.get("org_id"))
+                    print(
+                        f"{i}) {privilege.get('name')} (id: {privilege.get('org_id')})"
+                    )
                     orgs_with_sites.append(
                         {
-                            "org_id": privilege["org_id"],
-                            "name": privilege["name"],
+                            "org_id": privilege.get("org_id"),
+                            "name": privilege.get("name"),
                             "sites": [
                                 {
-                                    "site_id": privilege["site_id"],
-                                    "name": privilege["name"],
+                                    "site_id": privilege.get("site_id"),
+                                    "name": privilege.get("name"),
                                 }
                             ],
                         }
                     )
                 else:
                     orgs_with_sites[index]["sites"].append(
-                        {"site_id": privilege["site_id"], "name": privilege["name"]}
+                        {
+                            "site_id": privilege.get("site_id"),
+                            "name": privilege.get("name"),
+                        }
                     )
 
         if allow_many:
@@ -273,7 +283,7 @@ def select_site(
     if org_id is None:
         org_id = select_org(mist_session)[0]
 
-    for privilege in mist_session.privileges.privileges:
+    for privilege in mist_session.privileges:
         if privilege.get("scope") == "org" and privilege.get("org_id") == org_id:
             org_access = True
         if privilege.get("scope") == "site" and privilege.get("org_id") == org_id: