VDK-Plugin: Oauth plugin (#3410)

anitaratansingh · pre-commit-ci[bot] · web-flow · commit 7c2bd451320b · 2024-08-22T12:11:44.000+03:00
Added oauth plugin

---------

Co-authored-by: pre-commit-ci[bot] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
diff --git a/projects/vdk-core/src/vdk/internal/builtin_plugins/config/vdk_config.py b/projects/vdk-core/src/vdk/internal/builtin_plugins/config/vdk_config.py
@@ -162,19 +162,19 @@ def vdk_configure(self, config_builder: ConfigurationBuilder) -> None:
         )
         config_builder.add(
             TEAM_CLIENT_ID,
-            "",
+            None,
             True,
             "The Team's oAuth Client Id to use in authentication operations.",
         )
         config_builder.add(
             TEAM_CLIENT_SECRET,
-            "",
+            None,
             True,
             "The Team's oAuth Client Secret to use in authentication operations",
         )
         config_builder.add(
             TEAM_OAUTH_AUTHORIZE_URL,
-            "",
+            None,
             True,
             "The URL for Team's oAuth authorization",
         )
diff --git a/projects/vdk-plugins/vdk-kerberos-auth/src/vdk/plugin/kerberos/kerberos_configuration.py b/projects/vdk-plugins/vdk-kerberos-auth/src/vdk/plugin/kerberos/kerberos_configuration.py
@@ -18,8 +18,8 @@
 KEYTAB_REALM = "KEYTAB_REALM"
 KERBEROS_KDC_HOST = "KERBEROS_KDC_HOST"
 KRB_AUTH_FAIL_FAST = "KRB_AUTH_FAIL_FAST"
-
 API_SERVER_KERBEROS_SERVICE_NAME = "API_SERVER_KERBEROS_SERVICE_NAME"
+DISABLE_KERBEROS_LOGIN = "DISABLE_KERBEROS_LOGIN"
 
 
 class KerberosPluginConfiguration:
@@ -87,6 +87,9 @@ def auth_fail_fast(self) -> bool:
     def api_server_kerberos_service_name(self) -> str:
         return self.__config.get_value(API_SERVER_KERBEROS_SERVICE_NAME)
 
+    def disable_kerberos_plugin(self):
+        return self.__config.get_value(DISABLE_KERBEROS_LOGIN).lower() == "true"
+
 
 def add_definitions(config_builder: ConfigurationBuilder) -> None:
     config_builder.add(
@@ -152,3 +155,8 @@ def add_definitions(config_builder: ConfigurationBuilder) -> None:
     (for example, 'http@server.vdk.com').
     """,
     )
+    config_builder.add(
+        key=DISABLE_KERBEROS_LOGIN,
+        default_value=False,
+        description="To enable/disable kerberos login.",
+    )
diff --git a/projects/vdk-plugins/vdk-kerberos-auth/src/vdk/plugin/kerberos/kerberos_plugin.py b/projects/vdk-plugins/vdk-kerberos-auth/src/vdk/plugin/kerberos/kerberos_plugin.py
@@ -59,6 +59,8 @@ def vdk_initialize(self, context: CoreContext) -> None:
         kerberos_configuration = KerberosPluginConfiguration(
             None, None, context.configuration
         )
+        if kerberos_configuration.disable_kerberos_plugin():
+            return
         if (
             kerberos_configuration.keytab_filename()
             and kerberos_configuration.keytab_principal()
@@ -79,6 +81,9 @@ def initialize_job(self, context: JobContext) -> None:
         kerberos_configuration = KerberosPluginConfiguration(
             context.name, str(context.job_directory), context.core_context.configuration
         )
+        if kerberos_configuration.disable_kerberos_plugin():
+            return
+
         self.__attempt_kerberos_authentication(kerberos_configuration)
 
 
diff --git a/projects/vdk-plugins/vdk-kerberos-auth/tests/test_kerberos.py b/projects/vdk-plugins/vdk-kerberos-auth/tests/test_kerberos.py
@@ -41,6 +41,7 @@ def test_no_authentication(self):
             {
                 "VDK_KRB_AUTH_FAIL_FAST": "true",
                 "VDK_LOG_EXECUTION_RESULT": "True",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(
@@ -53,7 +54,11 @@ def test_no_authentication(self):
     def test_invalid_authentication_mode(self):
         with mock.patch.dict(
             os.environ,
-            {"VDK_KRB_AUTH": "invalid", "VDK_KRB_AUTH_FAIL_FAST": "true"},
+            {
+                "VDK_KRB_AUTH": "invalid",
+                "VDK_KRB_AUTH_FAIL_FAST": "true",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
+            },
         ):
             result: Result = self.__runner.invoke(
                 ["run", jobs_path_from_caller_directory("test-job")]
@@ -71,6 +76,7 @@ def test_kinit_authentication(self):
                 "VDK_KRB_AUTH_FAIL_FAST": "true",
                 "VDK_KRB5_CONF_FILENAME": krb5_conf_filename,
                 "VDK_LOG_EXECUTION_RESULT": "True",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(
@@ -86,6 +92,7 @@ def test_kinit_authentication_cli_command_no_auth(self):
             {
                 "VDK_KRB_AUTH": "",
                 "VDK_KRB_AUTH_FAIL_FAST": "true",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(["my-echo", "hi"])
@@ -106,6 +113,7 @@ def test_kinit_authentication_cli_command(self):
                 "VDK_KEYTAB_FOLDER": str(pathlib.Path(data_job_path).parent),
                 "VDK_KEYTAB_FILENAME": "test-job.keytab",
                 "VDK_KEYTAB_PRINCIPAL": "pa__view_test-job",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(["my-echo", "hi"])
@@ -125,6 +133,7 @@ def test_minikerberos_authentication_cli_command(self):
                 "VDK_KEYTAB_FOLDER": str(get_caller_directory().joinpath("jobs")),
                 "VDK_KEYTAB_FILENAME": "test-job.keytab",
                 "VDK_KEYTAB_PRINCIPAL": "pa__view_test-job",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(["my-echo", "hi"])
@@ -140,6 +149,7 @@ def test_minikerberos_authentication_cli_command_no_keytab_file(self):
                 "VDK_KRB_AUTH_FAIL_FAST": "true",
                 "VDK_KRB5_CONF_FILENAME": krb5_conf_filename,
                 "VDK_KEYTAB_PRINCIPAL": "pa__view_test-job",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(["my-echo", "hi"])
@@ -161,6 +171,7 @@ def test_kinit_authentication_with_wrong_credentials(self):
                         "different_principal.keytab"
                     )
                 ),
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(["run", data_job_path])
@@ -181,6 +192,7 @@ def test_kinit_authentication_error_fail_fast_is_false(self):
                 "VDK_KEYTAB_FILENAME": str(
                     pathlib.Path(data_job_path).parent.joinpath("non_existent.keytab")
                 ),
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(["run", data_job_path])
@@ -198,6 +210,7 @@ def test_kinit_authentication_with_missing_keytab(self):
                 "VDK_KRB5_CONF_FILENAME": krb5_conf_filename,
                 "VDK_KEYTAB_FOLDER": str(pathlib.Path(data_job_path).parent),
                 "VDK_KEYTAB_FILENAME": "non_existent.keytab",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(["run", data_job_path])
@@ -216,6 +229,7 @@ def test_minikerberos_authentication(self):
                 "VDK_KERBEROS_KDC_HOST": "localhost",
                 "VDK_KRB5_CONF_FILENAME": krb5_conf_filename,
                 "VDK_LOG_EXECUTION_RESULT": "True",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(
@@ -237,6 +251,7 @@ def test_minikerberos_authentication_with_wrong_credentials(self):
                 "VDK_KRB5_CONF_FILENAME": krb5_conf_filename,
                 "VDK_KEYTAB_FOLDER": str(pathlib.Path(data_job_path).parent),
                 "VDK_KEYTAB_FILENAME": "different_principal.keytab",
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(["run", data_job_path])
@@ -257,6 +272,7 @@ def test_minikerberos_authentication_with_missing_keytab(self):
                 "VDK_KEYTAB_FILENAME": str(
                     pathlib.Path(data_job_path).parent.joinpath("non_existent.keytab")
                 ),
+                "VDK_DISABLE_KERBEROS_LOGIN": "False",
             },
         ):
             result: Result = self.__runner.invoke(["run", data_job_path])
diff --git a/projects/vdk-plugins/vdk-oauth-auth/README.md b/projects/vdk-plugins/vdk-oauth-auth/README.md
@@ -0,0 +1,19 @@
+
+The plugin provides Oauth authentication on data job startup.
+
+# Usage
+
+To install the plugin, run:
+
+```bash
+pip install vdk-oauth-auth
+```
+
+## Configuration
+
+The following environment variables can be used to configure this plugin:
+
+| name                 | description                                   |
+|----------------------|-----------------------------------------------|
+| `TEAM_CLIENT_ID`     | Client id to fetch access token from CSP.     |
+| `TEAM_CLIENT_SECRET` | Client secret to fetch access token from CSP. |
diff --git a/projects/vdk-plugins/vdk-oauth-auth/requirements.txt b/projects/vdk-plugins/vdk-oauth-auth/requirements.txt
@@ -0,0 +1,8 @@
+docker<7
+docker-compose
+
+
+# install earlier version due to https://github.com/yaml/pyyaml/issues/601
+PyYAML==5.3.1
+
+vdk-test-utils
diff --git a/projects/vdk-plugins/vdk-oauth-auth/setup.py b/projects/vdk-plugins/vdk-oauth-auth/setup.py
@@ -0,0 +1,33 @@
+# Copyright 2023-2024 Broadcom
+# SPDX-License-Identifier: Apache-2.0
+import pathlib
+
+import setuptools
+
+
+__version__ = "0.3.0"
+
+setuptools.setup(
+    name="vdk-oauth-auth",
+    version=__version__,
+    url="https://github.com/vmware/versatile-data-kit",
+    description="Versatile Data Kit SDK plugin adds Oauth support.",
+    long_description=pathlib.Path("README.md").read_text(),
+    long_description_content_type="text/markdown",
+    install_requires=[
+        "vdk-core",
+    ],
+    package_dir={"": "src"},
+    packages=setuptools.find_namespace_packages(where="src"),
+    entry_points={"vdk.plugin.run": ["vdk-oauth-auth = vdk.plugin.oauth.oauth_plugin"]},
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+    ],
+)
diff --git a/projects/vdk-plugins/vdk-oauth-auth/src/vdk/plugin/oauth/oauth_configuration.py b/projects/vdk-plugins/vdk-oauth-auth/src/vdk/plugin/oauth/oauth_configuration.py
@@ -0,0 +1,55 @@
+# Copyright 2024-2025 Broadcom
+# SPDX-License-Identifier: Apache-2.0
+from vdk.internal.core.config import Configuration
+from vdk.internal.core.config import ConfigurationBuilder
+
+CLIENT_ID = "CLIENT_ID"
+CLIENT_SECRET = "CLIENT_SECRET"
+TEAM_CLIENT_ID = "TEAM_CLIENT_ID"
+TEAM_CLIENT_SECRET = "TEAM_CLIENT_SECRET"
+CONTROL_SERVICE_REST_API_URL = "CONTROL_SERVICE_REST_API_URL"
+API_TOKEN_AUTHORIZATION_URL = "API_TOKEN_AUTHORIZATION_URL"
+TEAM_OAUTH_AUTHORIZE_URL = "TEAM_OAUTH_AUTHORIZE_URL"
+CSP_ACCESS_TOKEN = "CSP_ACCESS_TOKEN"
+DISABLE_OAUTH_LOGIN = "DISABLE_OAUTH_LOGIN"
+Team = "TEAM"
+
+
+class OauthPluginConfiguration:
+    def __init__(
+        self,
+        config: Configuration,
+    ):
+        self.__config = config
+
+    def team(self):
+        return self.__config.get_value(Team)
+
+    def team_client_id(self):
+        return self.__config.get_value(TEAM_CLIENT_ID)
+
+    def team_client_secret(self):
+        return self.__config.get_value(TEAM_CLIENT_SECRET)
+
+    def control_service_rest_api_url(self):
+        return self.__config.get_value(CONTROL_SERVICE_REST_API_URL)
+
+    def api_token_authorization_url(self):
+        return self.__config.get_value(API_TOKEN_AUTHORIZATION_URL)
+
+    def team_oauth_authorize_url(self):
+        return self.__config.get_value(TEAM_OAUTH_AUTHORIZE_URL)
+
+    def csp_access_token(self):
+        return self.__config.get_value(CSP_ACCESS_TOKEN)
+
+    def disable_oauth_plugin(self):
+        return self.__config.get_value(DISABLE_OAUTH_LOGIN).lower() == "true"
+
+
+def add_definitions(config_builder: ConfigurationBuilder) -> None:
+    config_builder.add(
+        key=DISABLE_OAUTH_LOGIN,
+        default_value=False,
+        description="To enable/disable oauth login.",
+    )
diff --git a/projects/vdk-plugins/vdk-oauth-auth/src/vdk/plugin/oauth/oauth_plugin.py b/projects/vdk-plugins/vdk-oauth-auth/src/vdk/plugin/oauth/oauth_plugin.py

Original file line number	Diff line number	Diff line change
`@@ -162,19 +162,19 @@ def vdk_configure(self, config_builder: ConfigurationBuilder) -> None:`
`162`	`162`	`)`
`163`	`163`	`config_builder.add(`
`164`	`164`	`TEAM_CLIENT_ID,`
`165`		`- "",`
	`165`	`+ None,`
`166`	`166`	`True,`
`167`	`167`	`"The Team's oAuth Client Id to use in authentication operations.",`
`168`	`168`	`)`
`169`	`169`	`config_builder.add(`
`170`	`170`	`TEAM_CLIENT_SECRET,`
`171`		`- "",`
	`171`	`+ None,`
`172`	`172`	`True,`
`173`	`173`	`"The Team's oAuth Client Secret to use in authentication operations",`
`174`	`174`	`)`
`175`	`175`	`config_builder.add(`
`176`	`176`	`TEAM_OAUTH_AUTHORIZE_URL,`
`177`		`- "",`
	`177`	`+ None,`
`178`	`178`	`True,`
`179`	`179`	`"The URL for Team's oAuth authorization",`
`180`	`180`	`)`