create patch for superset-opa

labrenbe · labrenbe · commit 1a15b1b0bf88 · 2024-12-20T12:43:13.000+01:00
diff --git a/superset-opa-integration/TODO.md b/superset-opa-integration/TODO.md
@@ -3,7 +3,6 @@
 - Test with UIF + caching in OPA (+ document this)
 - Implement changes in operator and CRD
 - Documentation (how to write rego rules for this)
-- Create a patch for superset image with OPA integration behind feature flag
 - Tests
 - "Sync interval" mechanism in superset to improve latency and not spam OPA
 - Mount OPA service discovery configMap
@@ -14,3 +13,4 @@
 - OPA rego rules returning user-to-role mappings
 - Error handling in case OPA is not available or the API call returns a non 200 code or if the result data is garbage (e.g. the UIF source system is not available)
 - Make OPA address etc configurable via service discovery
+- Create a patch for superset image with OPA integration
diff --git a/superset-opa-integration/manager.py b/superset-opa-integration/manager.py
@@ -1,22 +1,19 @@
-from http.client import HTTPException
-import os
-from flask import g
+from flask import current_app, g
 from flask_appbuilder.security.sqla.models import (Role, User)
+from http.client import HTTPException
 from opa_client.opa import OpaClient
 from superset.security.manager import SupersetSecurityManager
 from superset import conf
 from typing import (Optional, List, Tuple)
 
 import logging
 
-# logger = logging.get_logger(__name__)
 class OpaSupersetSecurityManager(SupersetSecurityManager):
-
     def get_user_roles(self, user: Optional[User] = None) -> List[Role]:
         if not user:
             user = g.user
         
-        default_role = self.resolve_role(self.get_default_role())
+        default_role = self.resolve_role(current_app.config.get("AUTH_USER_REGISTRATION_ROLE"))
 
         opa_role_names = self.get_opa_user_roles(user.username)
         logging.info(f'OPA returned roles: {opa_role_names}')
@@ -35,20 +32,18 @@ def get_user_roles(self, user: Optional[User] = None) -> List[Role]:
 
     def get_opa_user_roles(self, username: str) -> set[str]:
         """
-        Queries an Open Policy Agent instance for the roles of a given user and returns a list of role names.
-        Returns an empty list if an exception during the connection to Open Policy Agent is encountered or if the query result
-        is not a list.
+        Queries an Open Policy Agent instance for the roles of a given user.
+        
+        :returns: A list of role names or an empty list if an exception during the connection to OPA
+        is encountered or if OPA didn't return a list.
         """
         host, port, tls = self.resolve_opa_endpoint()
-        print(host)
-        print(port)
-        print(tls)
         client = OpaClient(host = host, port = port, ssl = tls)
         try:
           response = client.query_rule(
                   input_data = {'username': username},
-                  package_path = 'superset',
-                  rule_name = 'user_roles')
+                  package_path = current_app.config.get('STACKABLE_OPA_PACKAGE'),
+                  rule_name = current_app.config.get('STACKABLE_OPA_RULE'))
         except HTTPException as e:
            logging.error(f'Encountered an exception while querying OPA:\n{e}')
            return []
@@ -58,12 +53,12 @@ def get_opa_user_roles(self, username: str) -> set[str]:
           logging.error(f'The OPA query didn\'t return a list: {response}')
           return []
         return roles
-    
+
 
     def resolve_opa_endpoint(self) -> Tuple[str, int, bool]:
-       opa_endpoint = os.getenv('STACKABLE_OPA_ENDPOINT')
-       [protocol, host, port] = opa_endpoint.split(":")
-       return host.lstrip('/'), int(port.rstrip('/')), protocol == 'https'
+      opa_endpoint = current_app.config.get('STACKABLE_OPA_ENDPOINT')
+      [protocol, host, port] = opa_endpoint.split(":")
+      return host.lstrip('/'), int(port.rstrip('/')), protocol == 'https'
     
 
     def resolve_role(self, role_name: str) -> Role:
@@ -72,7 +67,3 @@ def resolve_role(self, role_name: str) -> Role:
         logging.info(f'Creating role {role_name} as it doesn\'t already exist.')
         self.add_role(role_name)
       return self.find_role(role_name)
-         
-  
-    def get_default_role(self) -> str:
-      return conf["AUTH_USER_REGISTRATION_ROLE"] if conf["AUTH_USER_REGISTRATION_ROLE"] else "Public"
diff --git a/superset-opa-integration/superset-custom-opa.yaml b/superset-opa-integration/superset-custom-opa.yaml
@@ -8,7 +8,7 @@ metadata:
   }
 spec:
   image:
-    custom: docker.stackable.tech/sandbox/superset:4.0.2-stackable0.0.0-dev-opaV2
+    custom: docker.stackable.tech/stackable/superset:4.0.2-stackable0.0.0-dev-opaV6
     productVersion: 4.0.2
     pullPolicy: Never
   clusterConfig:
@@ -30,88 +30,13 @@ spec:
               configMapKeyRef:
                 key: OPA
                 name: simple-opa
+    envOverrides:         
+      AUTH_USER_REGISTRATION_ROLE: Gamma
     configOverrides:
       superset_config.py:
         EXPERIMENTAL_FILE_HEADER: |
-          from http.client import HTTPException
-          import os
-          from flask import g
-          from flask_appbuilder.security.sqla.models import (Role, User)
-          from opa_client.opa import OpaClient
-          from superset.security.manager import SupersetSecurityManager
-          from superset import conf
-          from typing import (Optional, Tuple, List)
-
-          import logging
-
-          # logger = logging.get_logger(__name__)
-          class OpaSupersetSecurityManager(SupersetSecurityManager):
-
-              def get_user_roles(self, user: Optional[User] = None) -> List[Role]:
-                  if not user:
-                      user = g.user
-                  
-                  default_role = self.resolve_role(self.get_default_role())
-
-                  opa_role_names = self.get_opa_user_roles(user.username)
-                  logging.info(f'OPA returned roles: {opa_role_names}')
-
-                  opa_roles = set(map(self.resolve_role, opa_role_names))
-                  # Ensure that in case of a bad or no reponse from OPA each user will have at least one role.
-                  opa_roles.add(default_role)
-                  
-                  if set(user.roles) != opa_roles:
-                    logging.info(f'Found diff in {user.roles} vs. {opa_roles}')
-                    user.roles = list(opa_roles)
-                    self.update_user(user)
-
-                  return user.roles
-              
-
-              def get_opa_user_roles(self, username: str) -> set[str]:
-                  """
-                  Queries an Open Policy Agent instance for the roles of a given user and returns a list of role names.
-                  Returns an empty list if an exception during the connection to Open Policy Agent is encountered or if the query result
-                  is not a list.
-                  """
-
-                  host, port, tls = self.resolve_opa_endpoint()
-                  client = OpaClient(host = host, port = port, ssl = tls)
-                  try:
-                    response = client.query_rule(
-                            input_data = {'username': username},
-                            package_path = 'superset',
-                            rule_name = 'user_roles')
-                  except HTTPException as e:
-                    logging.error(f'Encountered an exception while querying OPA:\n{e}')
-                    return []
-                  roles = response.get('result')
-                  # If OPA didn't return a result or if the result is not a list, return no roles.
-                  if roles is None or type(roles).__name__ != "list":
-                    logging.error(f'The OPA query didn\'t return a list: {response}')
-                    return []
-                  return roles
-
-
-              def resolve_opa_endpoint(self) -> Tuple[str, int, bool]:
-                opa_endpoint = os.getenv('STACKABLE_OPA_ENDPOINT')
-                [protocol, host, port] = opa_endpoint.split(":")
-                return host.lstrip('/'), int(port.rstrip('/')), protocol == 'https'
-                          
-
-              def resolve_role(self, role_name: str) -> Role:
-                role = self.find_role(role_name)
-                if role is None:
-                  logging.info(f'Creating role {role_name} as it doesn\'t already exist.')
-                  self.add_role(role_name)
-                return self.find_role(role_name)
-                  
-            
-              def get_default_role(self) -> str:
-                return conf["AUTH_USER_REGISTRATION_ROLE"] if conf["AUTH_USER_REGISTRATION_ROLE"] else "Public"
-        AUTH_USER_REGISTRATION_ROLE: Gamma
+          from superset.security.manager import OpaSupersetSecurityManager
         # Maybe also ENABLE_TEMPLATE_PROCESSING
-        CUSTOM_SECURITY_MANAGER: OpaSupersetSecurityManager
         FEATURE_FLAGS: |-
           {
             'ENABLE_JAVASCRIPT_CONTROLS': True,
@@ -143,3 +68,10 @@ spec:
           {
             False
           }
+        # TODO: Add these line with superset operator
+        CUSTOM_SECURITY_MANAGER: OpaSupersetSecurityManager
+        AUTH_USER_REGISTRATION_ROLE: os.getenv('AUTH_USER_REGISTRATION_ROLE', 'Public')
+        STACKABLE_OPA_ENDPOINT: os.getenv('STACKABLE_OPA_ENDPOINT')
+        STACKABLE_OPA_PACKAGE: |-
+          "superset"
+        STACKABLE_OPA_RULE: os.getenv('STACKABLE_OPA_RULE', 'user_roles')
diff --git a/superset/stackable/patches/4.0.2/001-opa-integration.patch b/superset/stackable/patches/4.0.2/001-opa-integration.patch
@@ -0,0 +1,88 @@
+diff --git a/superset/security/manager.py b/superset/security/manager.py
+index e5a32e97a..6971cf59a 100644
+--- a/superset/security/manager.py
++++ b/superset/security/manager.py
+@@ -21,7 +21,7 @@ import logging
+ import re
+ import time
+ from collections import defaultdict
+-from typing import Any, Callable, cast, NamedTuple, Optional, TYPE_CHECKING, Union
++from typing import Any, Callable, cast, List, NamedTuple, Optional, Tuple, TYPE_CHECKING, Union
+ 
+ from flask import current_app, Flask, g, Request
+ from flask_appbuilder import Model
+@@ -45,7 +45,9 @@ from flask_appbuilder.security.views import (
+ from flask_appbuilder.widgets import ListWidget
+ from flask_babel import lazy_gettext as _
+ from flask_login import AnonymousUserMixin, LoginManager
++from http.client import HTTPException
+ from jwt.api_jwt import _jwt_global_obj
++from opa_client.opa import OpaClient
+ from sqlalchemy import and_, inspect, or_
+ from sqlalchemy.engine.base import Connection
+ from sqlalchemy.orm import eagerload
+@@ -2465,3 +2467,64 @@ class SupersetSecurityManager(  # pylint: disable=too-many-public-methods
+         return current_app.config["AUTH_ROLE_ADMIN"] in [
+             role.name for role in self.get_user_roles()
+         ]
++
++
++class OpaSupersetSecurityManager(SupersetSecurityManager):
++    def get_user_roles(self, user: Optional[User] = None) -> List[Role]:
++        if not user:
++            user = g.user
++        
++        default_role = self.resolve_role(current_app.config.get("AUTH_USER_REGISTRATION_ROLE"))
++
++        opa_role_names = self.get_opa_user_roles(user.username)
++        logging.info(f'OPA returned roles: {opa_role_names}')
++
++        opa_roles = set(map(self.resolve_role, opa_role_names))
++        # Ensure that in case of a bad or no reponse from OPA each user will have at least one role.
++        opa_roles.add(default_role)
++        
++        if set(user.roles) != opa_roles:
++          logging.info(f'Found diff in {user.roles} vs. {opa_roles}')
++          user.roles = list(opa_roles)
++          self.update_user(user)
++
++        return user.roles
++    
++
++    def get_opa_user_roles(self, username: str) -> set[str]:
++        """
++        Queries an Open Policy Agent instance for the roles of a given user.
++        
++        :returns: A list of role names or an empty list if an exception during the connection to OPA
++        is encountered or if OPA didn't return a list.
++        """
++        host, port, tls = self.resolve_opa_endpoint()
++        client = OpaClient(host = host, port = port, ssl = tls)
++        try:
++          response = client.query_rule(
++                  input_data = {'username': username},
++                  package_path = current_app.config.get('STACKABLE_OPA_PACKAGE'),
++                  rule_name = current_app.config.get('STACKABLE_OPA_RULE'))
++        except HTTPException as e:
++           logging.error(f'Encountered an exception while querying OPA:\n{e}')
++           return []
++        roles = response.get('result')
++        # If OPA didn't return a result or if the result is not a list, return no roles.
++        if roles is None or type(roles).__name__ != "list":
++          logging.error(f'The OPA query didn\'t return a list: {response}')
++          return []
++        return roles
++
++
++    def resolve_opa_endpoint(self) -> Tuple[str, int, bool]:
++      opa_endpoint = current_app.config.get('STACKABLE_OPA_ENDPOINT')
++      [protocol, host, port] = opa_endpoint.split(":")
++      return host.lstrip('/'), int(port.rstrip('/')), protocol == 'https'
++    
++
++    def resolve_role(self, role_name: str) -> Role:
++      role = self.find_role(role_name)
++      if role is None:
++        logging.info(f'Creating role {role_name} as it doesn\'t already exist.')
++        self.add_role(role_name)
++      return self.find_role(role_name)
diff --git a/superset/stackable/patches/4.0.2/OpaSupersetSecurityManager.py b/superset/stackable/patches/4.0.2/OpaSupersetSecurityManager.py
