Refactor auth_client_manager_factory.py in function get_auth_client_m…

…anager

Signed-off-by: Theodor Mihalache <[email protected]>

sdk/python/feast/permissions/client/arrow_flight_auth_interceptor.py

@@ -2,7 +2,7 @@
 
 from feast.permissions.auth.auth_type import AuthType
 from feast.permissions.auth_model import AuthConfig
-from feast.permissions.client.auth_client_manager_factory import get_auth_token
+from feast.permissions.client.client_auth_token import get_auth_token
 
 
 class FlightBearerTokenInterceptor(fl.ClientMiddleware):

sdk/python/feast/permissions/client/auth_client_manager.py

@@ -1,8 +1,49 @@
+import os
 from abc import ABC, abstractmethod
 
+from feast.permissions.auth.auth_type import AuthType
+from feast.permissions.auth_model import (
+    AuthConfig,
+    KubernetesAuthConfig,
+    OidcClientAuthConfig,
+)
+
 
 class AuthenticationClientManager(ABC):
     @abstractmethod
     def get_token(self) -> str:
         """Retrieves the token based on the authentication type configuration"""
         pass
+
+
+class AuthenticationClientManagerFactory(ABC):
+    def __init__(self, auth_config: AuthConfig):
+        self.auth_config = auth_config
+
+    def get_auth_client_manager(self) -> AuthenticationClientManager:
+        from feast.permissions.client.intra_comm_authentication_client_manager import (
+            IntraCommAuthClientManager,
+        )
+        from feast.permissions.client.kubernetes_auth_client_manager import (
+            KubernetesAuthClientManager,
+        )
+        from feast.permissions.client.oidc_authentication_client_manager import (
+            OidcAuthClientManager,
+        )
+
+        intra_communication_base64 = os.getenv("INTRA_COMMUNICATION_BASE64")
+        if intra_communication_base64:
+            return IntraCommAuthClientManager(
+                self.auth_config, intra_communication_base64
+            )
+
+        if self.auth_config.type == AuthType.OIDC.value:
+            assert isinstance(self.auth_config, OidcClientAuthConfig)
+            return OidcAuthClientManager(self.auth_config)
+        elif self.auth_config.type == AuthType.KUBERNETES.value:
+            assert isinstance(self.auth_config, KubernetesAuthConfig)
+            return KubernetesAuthClientManager(self.auth_config)
+        else:
+            raise RuntimeError(
+                f"No Auth client manager implemented for the auth type:${self.auth_config.type}"
+            )

sdk/python/feast/permissions/client/client_auth_token.py

@@ -0,0 +1,14 @@
+from feast.permissions.auth_model import (
+    AuthConfig,
+)
+from feast.permissions.client.auth_client_manager import (
+    AuthenticationClientManagerFactory,
+)
+
+
+def get_auth_token(auth_config: AuthConfig) -> str:
+    return (
+        AuthenticationClientManagerFactory(auth_config)
+        .get_auth_client_manager()
+        .get_token()
+    )

sdk/python/feast/permissions/client/grpc_client_auth_interceptor.py

@@ -4,7 +4,7 @@
 
 from feast.errors import FeastError
 from feast.permissions.auth_model import AuthConfig
-from feast.permissions.client.auth_client_manager_factory import get_auth_token
+from feast.permissions.client.client_auth_token import get_auth_token
 
 logger = logging.getLogger(__name__)
 

sdk/python/feast/permissions/client/http_auth_requests_wrapper.py

@@ -5,7 +5,7 @@
 from feast.permissions.auth_model import (
     AuthConfig,
 )
-from feast.permissions.client.auth_client_manager_factory import get_auth_token
+from feast.permissions.client.client_auth_token import get_auth_token
 
 
 class AuthenticatedRequestsSession(Session):

sdk/python/feast/permissions/client/intra_comm_authentication_client_manager.py

@@ -0,0 +1,33 @@
+import logging
+
+import jwt
+
+from feast.permissions.auth.auth_type import AuthType
+from feast.permissions.auth_model import AuthConfig, OidcAuthConfig
+from feast.permissions.client.auth_client_manager import AuthenticationClientManager
+
+logger = logging.getLogger(__name__)
+
+
+class IntraCommAuthClientManager(AuthenticationClientManager):
+    def __init__(self, auth_config: AuthConfig, intra_communication_base64: str):
+        assert isinstance(auth_config, OidcAuthConfig)
+
+        self.auth_config = auth_config
+        self.intra_communication_base64 = intra_communication_base64
+
+    def get_token(self):
+        if self.auth_config.type == AuthType.OIDC.value:
+            payload = {
+                "preferred_username": f"{self.intra_communication_base64}",  # Subject claim
+            }
+        elif self.auth_config.type == AuthType.KUBERNETES.value:
+            payload = {
+                "sub": f":::{self.intra_communication_base64}",  # Subject claim
+            }
+        else:
+            raise RuntimeError(
+                f"No Auth client manager implemented for the auth type:${self.auth_config.type}"
+            )
+
+        return jwt.encode(payload, "")