datalens-tech · ovsds · Dec 20, 2024
diff --git a/app/dl_control_api/dl_control_api/app_factory.py b/app/dl_control_api/dl_control_api/app_factory.py
@@ -18,6 +18,8 @@
 from dl_api_lib.app_settings import (
     ControlApiAppSettingsOS,
     ControlApiAppTestingsSettings,
+    NativeAuthSettingsOS,
+    NullAuthSettingsOS,
     ZitadelAuthSettingsOS,
 )
 from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
@@ -90,14 +92,25 @@ def _setup_auth_middleware(
         testing_app_settings: Optional[ControlApiAppTestingsSettings] = None,
     ) -> AuthSetupResult:
         self._settings: ControlApiAppSettingsOS
+        settings = self._settings.AUTH
 
-        if self._settings.AUTH is None or self._settings.AUTH.type == "NONE":
+        if settings is None or isinstance(settings, NullAuthSettingsOS):
             return self._setup_auth_middleware_none(app=app, testing_app_settings=testing_app_settings)
 
-        if self._settings.AUTH.type == "ZITADEL":
-            return self._setup_auth_middleware_zitadel(app=app)
+        if isinstance(settings, ZitadelAuthSettingsOS):
+            return self._setup_auth_middleware_zitadel(
+                settings=settings,
+                cadata=get_multiple_root_certificates(
+                    self._settings.CA_FILE_PATH,
+                    *self._settings.EXTRA_CA_FILE_PATHS,
+                ),
+                app=app,
+            )
 
-        raise ValueError(f"Unknown auth type: {self._settings.AUTH.type}")
+        if isinstance(settings, NativeAuthSettingsOS):
+            return self._setup_auth_middleware_native(settings=settings, app=app)
+
+        raise ValueError(f"Unknown auth type: {settings.type}")
 
     def _setup_auth_middleware_none(
         self,
@@ -119,29 +132,25 @@ def _setup_auth_middleware_none(
 
     def _setup_auth_middleware_zitadel(
         self,
+        settings: ZitadelAuthSettingsOS,
+        cadata: bytes,
         app: flask.Flask,
     ) -> AuthSetupResult:
-        assert isinstance(self._settings.AUTH, ZitadelAuthSettingsOS)
-
         import httpx
 
         import dl_zitadel
 
-        ca_data = get_multiple_root_certificates(
-            self._settings.CA_FILE_PATH,
-            *self._settings.EXTRA_CA_FILE_PATHS,
+        httpx_client = httpx.Client(
+            verify=ssl.create_default_context(cadata=cadata.decode("ascii")),
         )
-
         zitadel_client = dl_zitadel.ZitadelSyncClient(
-            base_client=httpx.Client(
-                verify=ssl.create_default_context(cadata=ca_data.decode("ascii")),
-            ),
-            base_url=self._settings.AUTH.BASE_URL,
-            project_id=self._settings.AUTH.PROJECT_ID,
-            client_id=self._settings.AUTH.CLIENT_ID,
-            client_secret=self._settings.AUTH.CLIENT_SECRET,
-            app_client_id=self._settings.AUTH.APP_CLIENT_ID,
-            app_client_secret=self._settings.AUTH.APP_CLIENT_SECRET,
+            base_client=httpx_client,
+            base_url=settings.BASE_URL,
+            project_id=settings.PROJECT_ID,
+            client_id=settings.CLIENT_ID,
+            client_secret=settings.CLIENT_SECRET,
+            app_client_id=settings.APP_CLIENT_ID,
+            app_client_secret=settings.APP_CLIENT_SECRET,
         )
         token_storage = dl_zitadel.ZitadelSyncTokenStorage(
             client=zitadel_client,
@@ -154,3 +163,22 @@ def _setup_auth_middleware_zitadel(
         LOGGER.info("Zitadel auth setup complete")
 
         return AuthSetupResult(us_auth_mode=USAuthMode.regular)
+
+    def _setup_auth_middleware_native(
+        self,
+        settings: NativeAuthSettingsOS,
+        app: flask.Flask,
+    ) -> AuthSetupResult:
+        assert isinstance(settings, NativeAuthSettingsOS)
+
+        import dl_auth_native
+
+        dl_auth_native.FlaskMiddleware.from_settings(
+            settings=dl_auth_native.MiddlewareSettings(
+                decoder_key=settings.JWT_KEY,
+                decoder_algorithms=[settings.JWT_ALGORITHM],
+            )
+        ).set_up(app=app)
+        LOGGER.info("Native auth setup complete")
+
+        return AuthSetupResult(us_auth_mode=USAuthMode.regular)
diff --git a/app/dl_control_api/pyproject.toml b/app/dl_control_api/pyproject.toml
@@ -33,6 +33,7 @@ dl-cache-engine = {path = "../../lib/dl_cache_engine"}
 dl-configs = {path = "../../lib/dl_configs"}
 dl-constants = {path = "../../lib/dl_constants"}
 dl-zitadel = {path = "../../lib/dl_zitadel", extras = ["flask"]}
+dl-auth-native = {path = "../../lib/dl_auth_native", extras = ["flask"]}
 httpx = "^0.27.0"
 
 [build-system]

diff --git a/app/dl_data_api/dl_data_api/app_factory.py b/app/dl_data_api/dl_data_api/app_factory.py
@@ -16,6 +16,8 @@
 from dl_api_lib.app_settings import (
     AppSettings,
     DataApiAppSettingsOS,
+    NativeAuthSettingsOS,
+    NullAuthSettingsOS,
     ZitadelAuthSettingsOS,
 )
 from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
@@ -145,18 +147,24 @@ def set_up_environment(
         return result
 
     def _get_auth_middleware(self) -> Middleware:
-        if self._settings.AUTH is None or self._settings.AUTH.type == "NONE":
+        settings = self._settings.AUTH
+
+        if settings is None or isinstance(settings, NullAuthSettingsOS):
             return self._get_auth_middleware_none()
 
-        if self._settings.AUTH.type == "ZITADEL":
+        if isinstance(settings, ZitadelAuthSettingsOS):
             return self._get_auth_middleware_zitadel(
-                ca_data=get_multiple_root_certificates(
+                settings=settings,
+                cadata=get_multiple_root_certificates(
                     self._settings.CA_FILE_PATH,
                     *self._settings.EXTRA_CA_FILE_PATHS,
                 ),
             )
 
-        raise ValueError(f"Unknown auth type: {self._settings.AUTH.type}")
+        if isinstance(settings, NativeAuthSettingsOS):
+            return self._get_auth_middleware_native(settings)
+
+        raise ValueError(f"Unknown auth type: {settings.type}")
 
     def _get_auth_middleware_none(
         self,
@@ -168,30 +176,24 @@ def _get_auth_middleware_none(
 
     def _get_auth_middleware_zitadel(
         self,
-        ca_data: bytes,
+        settings: ZitadelAuthSettingsOS,
+        cadata: bytes,
     ) -> Middleware:
-        self._settings: DataApiAppSettingsOS
-        assert isinstance(self._settings.AUTH, ZitadelAuthSettingsOS)
-
         import httpx
 
         import dl_zitadel
 
-        ca_data = get_multiple_root_certificates(
-            self._settings.CA_FILE_PATH,
-            *self._settings.EXTRA_CA_FILE_PATHS,
+        httpx_client = httpx.AsyncClient(
+            verify=ssl.create_default_context(cadata=cadata.decode("ascii")),
         )
-
         zitadel_client = dl_zitadel.ZitadelAsyncClient(
-            base_client=httpx.AsyncClient(
-                verify=ssl.create_default_context(cadata=ca_data.decode("ascii")),
-            ),
-            base_url=self._settings.AUTH.BASE_URL,
-            project_id=self._settings.AUTH.PROJECT_ID,
-            client_id=self._settings.AUTH.CLIENT_ID,
-            client_secret=self._settings.AUTH.CLIENT_SECRET,
-            app_client_id=self._settings.AUTH.APP_CLIENT_ID,
-            app_client_secret=self._settings.AUTH.APP_CLIENT_SECRET,
+            base_client=httpx_client,
+            base_url=settings.BASE_URL,
+            project_id=settings.PROJECT_ID,
+            client_id=settings.CLIENT_ID,
+            client_secret=settings.CLIENT_SECRET,
+            app_client_id=settings.APP_CLIENT_ID,
+            app_client_secret=settings.APP_CLIENT_SECRET,
         )
         token_storage = dl_zitadel.ZitadelAsyncTokenStorage(
             client=zitadel_client,
@@ -200,5 +202,20 @@ def _get_auth_middleware_zitadel(
             client=zitadel_client,
             token_storage=token_storage,
         )
+        LOGGER.info("Zitadel auth middleware is set up")
+        return middleware.get_middleware()
 
+    def _get_auth_middleware_native(
+        self,
+        settings: NativeAuthSettingsOS,
+    ) -> Middleware:
+        import dl_auth_native
+
+        middleware = dl_auth_native.AioHTTPMiddleware.from_settings(
+            settings=dl_auth_native.MiddlewareSettings(
+                decoder_key=settings.JWT_KEY,
+                decoder_algorithms=[settings.JWT_ALGORITHM],
+            )
+        )
+        LOGGER.info("Native auth middleware is set up")
         return middleware.get_middleware()
diff --git a/app/dl_data_api/pyproject.toml b/app/dl_data_api/pyproject.toml
@@ -33,7 +33,8 @@ dl-api-commons = {path = "../../lib/dl_api_commons"}
 dl-cache-engine = {path = "../../lib/dl_cache_engine"}
 dl-configs = {path = "../../lib/dl_configs"}
 dl-constants = {path = "../../lib/dl_constants"}
-dl-zitadel = {path = "../../lib/dl_zitadel", extras = ["zitadel"]}
+dl-zitadel = {path = "../../lib/dl_zitadel", extras = ["aiohttp"]}
+dl-auth-native = {path = "../../lib/dl_auth_native", extras = ["aiohttp"]}
 httpx = "^0.27.0"
 
 

diff --git a/lib/dl_api_lib/dl_api_lib/app_settings.py b/lib/dl_api_lib/dl_api_lib/app_settings.py
@@ -339,6 +339,14 @@ class ZitadelAuthSettingsOS(BaseAuthSettingsOS):
 BaseAuthSettingsOS.register("ZITADEL", ZitadelAuthSettingsOS)
 
 
+class NativeAuthSettingsOS(BaseAuthSettingsOS):
+    JWT_KEY: str
+    JWT_ALGORITHM: str
+
+
+BaseAuthSettingsOS.register("NATIVE", NativeAuthSettingsOS)
+
+
 class AppSettingsOS(
     dl_settings.WithFallbackGetAttr,
     dl_settings.WithFallbackEnvSource,

diff --git a/lib/dl_auth_native/dl_auth_native/__init__.py b/lib/dl_auth_native/dl_auth_native/__init__.py
@@ -1,6 +1,7 @@
 from .middlewares import (
     AioHTTPMiddleware,
     FlaskMiddleware,
+    MiddlewareSettings,
 )
 from .token import (
     DecodeError,
@@ -14,6 +15,7 @@
 __all__ = [
     "AioHTTPMiddleware",
     "FlaskMiddleware",
+    "MiddlewareSettings",
     "Decoder",
     "DecoderProtocol",
     "Payload",

diff --git a/lib/dl_auth_native/dl_auth_native/middlewares/__init__.py b/lib/dl_auth_native/dl_auth_native/middlewares/__init__.py
@@ -1,8 +1,10 @@
 from .aiohttp import AioHTTPMiddleware
+from .base import MiddlewareSettings
 from .flask import FlaskMiddleware
 
 
 __all__ = [
     "AioHTTPMiddleware",
     "FlaskMiddleware",
+    "MiddlewareSettings",
 ]
diff --git a/lib/dl_auth_native/dl_auth_native/middlewares/base.py b/lib/dl_auth_native/dl_auth_native/middlewares/base.py
@@ -1,6 +1,7 @@
 import typing
 
 import attr
+import typing_extensions
 
 import dl_api_commons.base_models as dl_api_commons_base_models
 import dl_auth_native.token as token
@@ -34,12 +35,29 @@ class AuthResult:
     auth_data: AuthData = attr.ib()
 
 
+@attr.s(frozen=True)
+class MiddlewareSettings:
+    decoder_key: str = attr.ib()
+    decoder_algorithms: list[str] = attr.ib()
+
+
 @attr.s()
 class BaseMiddleware:
     _token_decoder: token.DecoderProtocol = attr.ib()
     _user_access_header_key: str = attr.ib(default=dl_api_commons_base_models.DLHeadersCommon.AUTHORIZATION_TOKEN)
     _token_type: str = attr.ib(default="Bearer")
 
+    @classmethod
+    def from_settings(cls, settings: MiddlewareSettings) -> typing_extensions.Self:
+        token_decoder = token.Decoder(
+            key=settings.decoder_key,
+            algorithms=settings.decoder_algorithms,
+        )
+
+        return cls(
+            token_decoder=token_decoder,
+        )
+
     @attr.s(frozen=True)
     class Unauthorized(Exception):
         message: str = attr.ib()
@@ -72,4 +90,5 @@ def _auth(self, user_access_header: typing.Optional[str]) -> AuthResult:
     "BaseMiddleware",
     "AuthResult",
     "AuthData",
+    "MiddlewareSettings",
 ]
diff --git a/lib/dl_auth_native/pyproject.toml b/lib/dl_auth_native/pyproject.toml
@@ -10,16 +10,21 @@ readme = "README.md"
 
 
 [tool.poetry.dependencies]
-aiohttp = ">=3.9.1"
+aiohttp = {version = ">=3.9.1", optional = true}
 attrs = ">=22.2.0"
-flask = ">=2.2.5"
+flask = {version = ">=2.2.5", optional = true}
 pyjwt = ">=2.4.0"
 pydantic = ">=2.7.0"
 python = ">=3.10, <3.13"
+typing-extensions = ">=4.9.0"
 werkzeug = ">=2.2.3"
 dl-constants = {path = "../../lib/dl_constants"}
 dl-api-commons = {path = "../../lib/dl_api_commons"}
 
+[tool.poetry.extras]
+aiohttp = ["aiohttp"]
+flask = ["flask"]
+
 [tool.poetry.group.tests.dependencies]
 pytest = ">=7.2.2"
 cryptography = ">=41.0.4"