diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml index 0932c8a95eaf2..95ce1b0acce36 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml +++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml @@ -85,6 +85,43 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + /auth/token/cli: + post: + tags: + - KeycloakAuthManagerToken + summary: Create Token Cli + operationId: create_token_cli + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/TokenBody' + required: true + responses: + '201': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/TokenResponse' + '400': + description: Bad Request + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '401': + description: Unauthorized + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' components: schemas: HTTPExceptionResponse: diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/token.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/token.py index e40434f6ff2c1..b3bb1bc1ff331 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/token.py +++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/token.py @@ -19,16 +19,13 @@ import logging -from fastapi import HTTPException -from keycloak import KeycloakAuthenticationError from starlette import status -from airflow.api_fastapi.app import get_auth_manager from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.configuration import conf from airflow.providers.keycloak.auth_manager.datamodels.token import TokenBody, TokenResponse -from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import KeycloakAuthManager -from airflow.providers.keycloak.auth_manager.user import KeycloakAuthManagerUser +from airflow.providers.keycloak.auth_manager.services.token import create_token_for log = logging.getLogger(__name__) token_router = AirflowRouter(tags=["KeycloakAuthManagerToken"]) @@ -40,23 +37,19 @@ responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST, status.HTTP_401_UNAUTHORIZED]), ) def create_token(body: TokenBody) -> TokenResponse: - client = KeycloakAuthManager.get_keycloak_client() - - try: - tokens = client.token(body.username, body.password) - except KeycloakAuthenticationError: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail="Invalid credentials", - ) - - userinfo = client.userinfo(tokens["access_token"]) - user = KeycloakAuthManagerUser( - user_id=userinfo["sub"], - name=userinfo["preferred_username"], - access_token=tokens["access_token"], - refresh_token=tokens["refresh_token"], - ) - token = get_auth_manager().generate_jwt(user) + token = create_token_for(body.username, body.password) + return TokenResponse(access_token=token) + +@token_router.post( + "/token/cli", + status_code=status.HTTP_201_CREATED, + responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST, status.HTTP_401_UNAUTHORIZED]), +) +def create_token_cli(body: TokenBody) -> TokenResponse: + token = create_token_for( + body.username, + body.password, + expiration_time_in_seconds=int(conf.getint("api_auth", "jwt_cli_expiration_time")), + ) return TokenResponse(access_token=token) diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/__init__.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/token.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/token.py new file mode 100644 index 0000000000000..754966d7520ee --- /dev/null +++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/token.py @@ -0,0 +1,53 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from fastapi import HTTPException +from keycloak import KeycloakAuthenticationError +from starlette import status + +from airflow.api_fastapi.app import get_auth_manager +from airflow.configuration import conf +from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import KeycloakAuthManager +from airflow.providers.keycloak.auth_manager.user import KeycloakAuthManagerUser + + +def create_token_for( + username: str, + password: str, + expiration_time_in_seconds: int = conf.getint("api_auth", "jwt_expiration_time"), +) -> str: + client = KeycloakAuthManager.get_keycloak_client() + + try: + tokens = client.token(username, password) + except KeycloakAuthenticationError: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid credentials", + ) + + userinfo = client.userinfo(tokens["access_token"]) + user = KeycloakAuthManagerUser( + user_id=userinfo["sub"], + name=userinfo["preferred_username"], + access_token=tokens["access_token"], + refresh_token=tokens["refresh_token"], + ) + + return get_auth_manager().generate_jwt(user, expiration_time_in_seconds=expiration_time_in_seconds) diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py b/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py index 8e45100b26c56..b05b45ff485ae 100644 --- a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py +++ b/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py @@ -16,41 +16,46 @@ # under the License. from __future__ import annotations -from unittest.mock import Mock, patch - -from keycloak import KeycloakAuthenticationError +from unittest.mock import patch from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX +from tests_common.test_utils.config import conf_vars + class TestTokenRouter: - @patch("airflow.providers.keycloak.auth_manager.routes.login.KeycloakAuthManager.get_keycloak_client") - def test_create_token(self, mock_get_keycloak_client, client): - mock_keycloak_client = Mock() - mock_keycloak_client.token.return_value = { - "access_token": "access_token", - "refresh_token": "refresh_token", + token = "token" + token_body_dict = {"username": "username", "password": "password"} + + @conf_vars( + { + ("api_auth", "jwt_expiration_time"): "10", } - mock_keycloak_client.userinfo.return_value = {"sub": "sub", "preferred_username": "username"} - mock_get_keycloak_client.return_value = mock_keycloak_client + ) + @patch("airflow.providers.keycloak.auth_manager.routes.token.create_token_for") + def test_create_token(self, mock_create_token_for, client): + mock_create_token_for.return_value = self.token response = client.post( AUTH_MANAGER_FASTAPI_APP_PREFIX + "/token", - json={"username": "username", "password": "password"}, + json=self.token_body_dict, ) assert response.status_code == 201 - mock_keycloak_client.token.assert_called_once_with("username", "password") - mock_keycloak_client.userinfo.assert_called_once_with("access_token") - - @patch("airflow.providers.keycloak.auth_manager.routes.login.KeycloakAuthManager.get_keycloak_client") - def test_create_token_with_invalid_creds(self, mock_get_keycloak_client, client): - mock_keycloak_client = Mock() - mock_keycloak_client.token.side_effect = KeycloakAuthenticationError() - mock_get_keycloak_client.return_value = mock_keycloak_client + assert response.json() == {"access_token": self.token} + + @conf_vars( + { + ("api_auth", "jwt_cli_expiration_time"): "10", + ("api_auth", "jwt_expiration_time"): "10", + } + ) + @patch("airflow.providers.keycloak.auth_manager.routes.token.create_token_for") + def test_create_token_cli(self, mock_create_token_for, client): + mock_create_token_for.return_value = self.token response = client.post( - AUTH_MANAGER_FASTAPI_APP_PREFIX + "/token", - json={"username": "username", "password": "password"}, + AUTH_MANAGER_FASTAPI_APP_PREFIX + "/token/cli", + json=self.token_body_dict, ) - assert response.status_code == 401 - mock_keycloak_client.token.assert_called_once_with("username", "password") + assert response.status_code == 201 + assert response.json() == {"access_token": self.token} diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/services/__init__.py b/providers/keycloak/tests/unit/keycloak/auth_manager/services/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/keycloak/tests/unit/keycloak/auth_manager/services/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/services/test_token.py b/providers/keycloak/tests/unit/keycloak/auth_manager/services/test_token.py new file mode 100644 index 0000000000000..7fe47c08474d3 --- /dev/null +++ b/providers/keycloak/tests/unit/keycloak/auth_manager/services/test_token.py @@ -0,0 +1,78 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest.mock import Mock, patch + +import fastapi +import pytest +from keycloak import KeycloakAuthenticationError + +from airflow.configuration import conf +from airflow.providers.keycloak.auth_manager.services.token import create_token_for + +from tests_common.test_utils.config import conf_vars + + +class TestTokenService: + token = "token" + test_username = "test_user" + test_password = "test_pass" + test_access_token = "access_token" + test_refresh_token = "refresh_token" + + @conf_vars( + { + ("api_auth", "jwt_expiration_time"): "10", + } + ) + @patch("airflow.providers.keycloak.auth_manager.services.token.get_auth_manager") + @patch("airflow.providers.keycloak.auth_manager.services.token.KeycloakAuthManager.get_keycloak_client") + def test_create_token(self, mock_get_keycloak_client, mock_get_auth_manager): + mock_keycloak_client = Mock() + mock_keycloak_client.token.return_value = { + "access_token": self.test_access_token, + "refresh_token": self.test_refresh_token, + } + mock_keycloak_client.userinfo.return_value = {"sub": "sub", "preferred_username": "username"} + mock_get_keycloak_client.return_value = mock_keycloak_client + mock_auth_manager = Mock() + mock_get_auth_manager.return_value = mock_auth_manager + mock_auth_manager.generate_jwt.return_value = self.token + + assert create_token_for(username=self.test_username, password=self.test_password) == self.token + mock_keycloak_client.token.assert_called_once_with(self.test_username, self.test_password) + mock_keycloak_client.userinfo.assert_called_once_with(self.test_access_token) + + @conf_vars( + { + ("api_auth", "jwt_cli_expiration_time"): "10", + ("api_auth", "jwt_expiration_time"): "10", + } + ) + @patch("airflow.providers.keycloak.auth_manager.services.token.KeycloakAuthManager.get_keycloak_client") + def test_create_token_with_invalid_creds(self, mock_get_keycloak_client): + mock_keycloak_client = Mock() + mock_keycloak_client.token.side_effect = KeycloakAuthenticationError() + mock_get_keycloak_client.return_value = mock_keycloak_client + + with pytest.raises(fastapi.exceptions.HTTPException): + create_token_for( + username=self.test_username, + password=self.test_password, + expiration_time_in_seconds=conf.getint("api_auth", "jwt_cli_expiration_time"), + )