Merge pull request apache#111 from riskive/ZFE-85176-internal-rbac

Zfe 85176 internal rbac

bi_superset/bi_cli/bi_cli.py

@@ -38,4 +38,9 @@ def bi_init() -> None:
     sec_manager.update_dashboard_default_access()
     logging.info("Updating Superset Dashboard default access Completed")
 
+    logging.info("Updating RBAC Roles")
+    sec_manager.sync_rbac_role_list()
+    sec_manager.sync_dashboard_rbac_role_assignation()
+    logging.info("Updating RBAC Roles Completed")
+
     logging.info("Sync Completed")

bi_superset/bi_cli/bi_cli_security_manager.py

@@ -191,6 +191,7 @@ def loads_roles_per_job_title(self):
                 dtype={
                     "username": String(255),
                     "role_name": String(255),
+                    "rbac_roles": String(255),
                 },
                 index=False,
             )
@@ -302,3 +303,165 @@ def update_dashboard_default_access(self):
             session.merge(dashboard)
         # Commit changes
         session.commit()
+
+    def sync_rbac_role_list(self):
+        if not AccessMethod.is_internal(self._access_method):
+            logging.info("This method only works with internal access method")
+            return
+
+        from bi_superset.bi_security_manager.adapter.bigquery_sql import BigquerySQL
+        from bi_superset.bi_security_manager.services.dashboard_role_access_service import (
+            DashboardRoleAccessService,
+        )
+        from superset.utils.database import get_main_database
+        from sqlalchemy import inspect, String, Integer
+        from bi_superset.bi_security_manager.models.models import (
+                            RBACRoles,
+                        )
+        from superset.models.dashboard import DashboardRoles
+
+        from flask_appbuilder.security.sqla.models import Role, assoc_user_role
+
+        bq_client = BigquerySQL()
+        dashabord_role_access_service = DashboardRoleAccessService(bq_client)
+
+        roles_per_job_title_df = dashabord_role_access_service.get_rbac_roles()
+
+        tbl_name = "bi_rbac_roles"
+        database = get_main_database()
+        with database.get_sqla_engine_with_context() as engine:
+            schema = inspect(engine).default_schema_name
+
+            roles_per_job_title_df.to_sql(
+                tbl_name,
+                engine,
+                schema=schema,
+                if_exists="replace",
+                chunksize=500,
+                dtype={
+                    "id": Integer,
+                    "role_name": String(255),
+                },
+                index_label="id",
+            )
+        session = self.get_session()
+
+
+        logging.info("Gettint RBAC Roles from Role table")
+        query = (
+                    session
+                    .query(Role)
+                    .filter(Role.name.startswith("rbac_"))
+                )
+        roles = query.all()
+
+        logging.info("Getting RBAC Roles from RBAC Role Table")
+        query = (
+                    session
+                    .query(RBACRoles)
+                    .filter(RBACRoles.role_name.startswith("rbac_"))
+                )
+        rbac_roles = query.all()
+
+        # delete roles thar are not in rbac_roles
+        for role in roles:
+            # Search by role.name in rbac_roles.role_name
+            if role.name not in [rbac_role.role_name for rbac_role in rbac_roles]:
+                logging.info("Deleting role from superset %s", role.name)
+                # need to search all role_id in `assoc_user_role` and delete
+                # Delete user role association
+                session.query(assoc_user_role).filter(assoc_user_role.c.role_id == role.id).delete(synchronize_session=False)
+                session.commit()
+                # Delete dashboard role association
+                session.query(DashboardRoles).filter(DashboardRoles.c.role_id == role.id).delete(synchronize_session=False)
+                session.commit()
+                # Delete Role
+                session.delete(role)
+                session.commit()
+
+        # Skip existing RBAC Roles
+        for rbac_role in rbac_roles:
+            if rbac_role.role_name not in [role.name for role in roles]:
+                logging.info("Creating role from superset %s", rbac_role.role_name)
+                self.add_role(rbac_role.role_name)
+
+
+    def sync_dashboard_rbac_role_assignation(self):
+        if not AccessMethod.is_internal(self._access_method):
+            logging.info("This method only works with internal access method")
+            return
+
+        from bi_superset.bi_security_manager.adapter.bigquery_sql import BigquerySQL
+        from bi_superset.bi_security_manager.services.dashboard_role_access_service import (
+            DashboardRoleAccessService,
+        )
+        from superset.utils.database import get_main_database
+        from sqlalchemy import inspect, String, Integer
+
+
+
+        bq_client = BigquerySQL()
+        dashabord_role_access_service = DashboardRoleAccessService(bq_client)
+
+        roles_per_job_title_df = dashabord_role_access_service.get_dashboard_rbac_role_assignation()
+
+        tbl_name = "bi_dashboard_rbac_role_assignation"
+        database = get_main_database()
+        with database.get_sqla_engine_with_context() as engine:
+            schema = inspect(engine).default_schema_name
+
+            roles_per_job_title_df.to_sql(
+                tbl_name,
+                engine,
+                schema=schema,
+                if_exists="replace",
+                chunksize=500,
+                dtype={
+                    "id": Integer,
+                    "role_name": String(255),
+                    "dashboard_id": Integer,
+                },
+                index_label="id",
+            )
+
+        from superset.models.dashboard import Dashboard
+        session = self.get_session()
+
+        # Get All Dashboards
+        query = session.query(Dashboard)
+
+        dashboards = query.all()
+
+
+
+        from bi_superset.bi_security_manager.models.models import (
+                            DashboardRBACRoleAssignation,
+                        )
+        import datetime
+        query = (
+                    session
+                    .query(DashboardRBACRoleAssignation)
+                    .filter(DashboardRBACRoleAssignation.role_name.startswith("rbac_"))
+                )
+        rbac_roles = query.all()
+
+        # Every Dashboard will remove role relationship
+        # where dashboard_id is not -1
+        exclusion_dashboard_ids = [rbac_role.dashboard_id for rbac_role in rbac_roles if rbac_role.dashboard_id != -1]
+
+        for dashboard in dashboards:
+            roles =  []
+            for rbac_role in rbac_roles:
+                if rbac_role.dashboard_id == dashboard.id:
+                    role = self.find_role(rbac_role.role_name)
+                    if role is None:
+                        raise Exception("Role not found")
+                    roles.append(role)
+                elif rbac_role.dashboard_id == -1 and dashboard.id not in exclusion_dashboard_ids:
+                    role = self.find_role(rbac_role.role_name)
+                    if role is None:
+                        raise Exception("Role not found")
+                    roles.append(role)
+            dashboard.roles = roles
+            session.merge(dashboard)
+        session.commit()

bi_superset/bi_security_manager/models/models.py

@@ -9,6 +9,7 @@ class RolesPerJobTitle(Model):
 
     username = Column(String(256), primary_key=True)
     role_name = Column(String(256))
+    rbac_roles = Column(String(256))
 
     def __repr__(self):
         return self.username
@@ -18,12 +19,20 @@ def from_dict(cls, i_dict) -> "RolesPerJobTitle":
         return cls(
             username=i_dict.get("employee"),
             role_name=i_dict.get("role_name").lower().replace(" ", "_"),
+            rbac_roles=i_dict.get("rbac_roles").lower() if i_dict.get("rbac_roles") else None,
         )
+
+    def list_rbac_roles(self):
+        # lower and add 'rbac_'
+        data = self.rbac_roles.split(",") if self.rbac_roles else []
+        return [f"rbac_{role.lower().strip().replace(' ', '_')}" for role in data]
+
 
     def to_dict(self) -> dict:
         return {
             "username": self.username,
             "role_name": self.role_name,
+            "rbac_roles": self.rbac_roles,
         }
 
 
@@ -75,3 +84,40 @@ def to_dict(self):
             "role_name": self.role_name,
             "dashboard_id": self.dashboard_id,
         }
+
+class RBACRoles(Model):
+    __tablename__ = "bi_rbac_roles"
+
+    id = Column(Integer, Sequence("dashboard_rbac_id_seq"), primary_key=True)
+    role_name = Column(String(256))
+
+    @classmethod
+    def from_dict(cls, row):
+        return cls(
+            role_name=row["role_name"].strip().lower().replace(" ", "_")
+        )
+
+    def to_dict(self):
+        return {
+            "role_name": self.role_name
+        }
+
+class DashboardRBACRoleAssignation(Model):
+    __tablename__ = "bi_dashboard_rbac_role_assignation"
+
+    id = Column(Integer, Sequence("dashboard_rbac_id_seq"), primary_key=True)
+    dashboard_id = Column(Integer)
+    role_name = Column(String(256))
+
+    @classmethod
+    def from_dict(cls, row):
+        return cls(
+            role_name=row["role_name"].strip().lower().replace(" ", "_"),
+            dashboard_id=row["dashboard_id"],
+        )
+
+    def to_dict(self):
+        return {
+            "role_name": self.role_name,
+            "dashboard_id": self.dashboard_id
+        }

bi_superset/bi_security_manager/services/dashboard_role_access_service.py

@@ -1,9 +1,11 @@
 import pandas as pd
 
 from bi_superset.bi_security_manager.port.a_sql import ASql
-from bi_superset.bi_security_manager.sql.queries import DASHBOARD_ROLE_ACCESS_EXTERNAL
+from bi_superset.bi_security_manager.sql.queries import DASHBOARD_ROLE_ACCESS_EXTERNAL, RBAC_ROLES, DASHBOARD_RBAC_ROLE_ASSIGNATION
 from bi_superset.bi_security_manager.models.models import (
     DashboardRoleAccessExternal,
+    RBACRoles,
+    DashboardRBACRoleAssignation
 )
 from flask import current_app
 
@@ -14,6 +16,7 @@ class DashboardRoleAccessService:
     def __init__(self, sql: ASql):
         self.sql: ASql = sql
 
+
     def get_dashboard_role_access_external(self):
         """
         Get dashboard role access for external users
@@ -39,3 +42,48 @@ def get_dashboard_role_access_external(self):
         )
         # nmapping
         return df
+
+    def get_rbac_roles(self):
+        """Retrieves RBAC roles from the database
+
+        Returns:
+            pd.DataFrame: Return a Dataframe of the RBAC roles
+        """
+        query = RBAC_ROLES.format(dataset=BQ_DATASET)
+        df = self.sql.get_df(query)
+        if df.empty:
+            raise ValueError("no rbac roles found")
+
+        rbac_roles = []
+        for _, row in df.iterrows():
+            rbac_role = RBACRoles.from_dict(row)
+            rbac_roles.append(rbac_role)
+
+        df = pd.DataFrame(
+            [
+                rbac_role.to_dict()
+                for rbac_role in rbac_roles
+            ]
+        )
+        # nmapping
+        return df
+
+    def get_dashboard_rbac_role_assignation(self):
+        query = DASHBOARD_RBAC_ROLE_ASSIGNATION.format(dataset=BQ_DATASET)
+        df = self.sql.get_df(query)
+        if df.empty:
+            raise ValueError("no rbac role assignation found")
+
+        dashboard_rbac_role_assignations = []
+        for _, row in df.iterrows():
+            dashboard_rbac_role_assignation = DashboardRBACRoleAssignation.from_dict(row)
+            dashboard_rbac_role_assignations.append(dashboard_rbac_role_assignation)
+
+        df = pd.DataFrame(
+            [
+                dashboard_rbac_role_assignation.to_dict()
+                for dashboard_rbac_role_assignation in dashboard_rbac_role_assignations
+            ]
+        )
+        # nmapping
+        return df

bi_superset/bi_security_manager/services/user_service.py

@@ -73,7 +73,7 @@ def get_and_update_user_roles(self, user, zf_user: ZFUser):
             self._access_origin
         ):
             if AccessMethod.is_external(self._access_method):
-                role = self.sm.find_role("Admin")
+                roles = [self.sm.find_role("Admin")]
             else:
                 # pylint: disable=import-outside-toplevel
                 from bi_superset.bi_security_manager.models.models import (
@@ -92,11 +92,19 @@ def get_and_update_user_roles(self, user, zf_user: ZFUser):
                 search_role_name = user_role_job_title.role_name
                 if "admin" == search_role_name.lower():
                     search_role_name = search_role_name.capitalize()
+                roles = []
+
                 role = self.sm.find_role(search_role_name)
                 if role is None:
                     logger.info("Role Not found")
+                roles.append(role)
                 # role = self.sm.find_role("zerofox_internal")
-            user.roles = [role]
+                rbac_roles = user_role_job_title.list_rbac_roles()
+                for rbac_role in rbac_roles:
+                    role = self.sm.find_role(rbac_role)
+                    if role is not None:
+                        roles.append(role)
+            user.roles = roles
 
         else:
             user.roles = self._get_user_roles(zf_user)

bi_superset/bi_security_manager/sql/queries.py

@@ -27,7 +27,8 @@
 ROLES_PER_JOB_TITLE = """
 SELECT
     employee,
-    role_name
+    role_name,
+    rbac_roles
 FROM {dataset}.bi_superset_access.roles_per_job_title
 """
 
@@ -52,3 +53,24 @@
     role_name
 FROM {dataset}.bi_superset_access.dashboard_role_access_exterrnal
 """
+
+##################################
+# DASHBOARD RBAC INTERNAL        #
+##################################
+
+RBAC_ROLES = """
+SELECT
+    concat('rbac_', role_name) as role_name
+FROM {dataset}.bi_superset_access.rbac_roles
+"""
+
+###################################
+# DASHBOARD RBAC ROLE ASSIGNATION #
+###################################
+
+DASHBOARD_RBAC_ROLE_ASSIGNATION = """
+SELECT
+    concat('rbac_', role_name) as role_name,
+    dashboard_id
+FROM {dataset}.bi_superset_access.dashboard_rbac_role_assignation
+"""