diff --git a/airflow-ctl/docs/images/command_hashes.txt b/airflow-ctl/docs/images/command_hashes.txt index 830500bf60b69..30ebffe3e6975 100644 --- a/airflow-ctl/docs/images/command_hashes.txt +++ b/airflow-ctl/docs/images/command_hashes.txt @@ -14,4 +14,4 @@ tasks:eb70701dfe1b9baeda39d5eea07fde5f variables:f8fc76d3d398b2780f4e97f7cd816646 version:31f4efdf8de0dbaaa4fac71ff7efecc3 plugins:4864fd8f356704bd2b3cd1aec3567e35 -auth login:963355e38ab1c133f2cf9495d00d99f6 +auth login:e068cc2df6160c58b93ef21b3d51abca diff --git a/airflow-ctl/docs/images/output_auth_login.svg b/airflow-ctl/docs/images/output_auth_login.svg index f722fb20eb3d3..ebc2556c4921d 100644 --- a/airflow-ctl/docs/images/output_auth_login.svg +++ b/airflow-ctl/docs/images/output_auth_login.svg @@ -1,4 +1,4 @@ - + - - + + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + + + + + + + + + + + + + + + + + + + + + + + + + + + + - + - + - - Usage:airflowctl auth login [-h] [--api-tokenAPI_TOKEN] -                             [--api-urlAPI_URL] [-eENV] -                             [--passwordPASSWORD] [--skip-keyring] -                             [--usernameUSERNAME] - -Log in to the Airflow API server and store the credentials for the selected -environment. Pass --api-token (or set AIRFLOW_CLI_TOKEN), or --username and ---password; missing credentials are prompted for when running in a  -terminal. - -Options: --h, --helpshow this help message and exit ---api-tokenAPI_TOKEN -The token to use for authentication ---api-urlAPI_URLThe URL of the Airflow API server --e, --envENVThe environment to run the command in ---passwordPASSWORDThe password to use for authentication ---skip-keyringSkip storing credentials in keyring ---usernameUSERNAMEThe username to use for authentication - + + Usage:airflowctl auth login [-h] [--api-tokenAPI_TOKEN] +                             [--api-urlAPI_URL] [-eENV] [--no-browser] +                             [--passwordPASSWORD] [--skip-keyring] +                             [--sso] [--timeoutTIMEOUT] +                             [--usernameUSERNAME] + +Log in to the Airflow API server and store the credentials for the selected +environment. Pass --api-token (or set AIRFLOW_CLI_TOKEN), or --username and +--password; missing credentials are prompted for when running in a  +terminal. Pass --sso to log in through the browser when the API server  +offers it. + +Options: +-h, --helpshow this help message and exit +--api-tokenAPI_TOKEN +The token to use for authentication +--api-urlAPI_URLThe URL of the Airflow API server +-e, --envENVThe environment to run the command in +--no-browserWith --sso, print the login URL instead of opening  +a browser. The browser must still reach this machine on 127.0.0.1, for  +example through 'ssh -L' when logged in remotely +--passwordPASSWORDThe password to use for authentication +--skip-keyringSkip storing credentials in keyring +--ssoLog in through the browser with the identity  +provider of the API server, if it offers it +--timeoutTIMEOUTWith --sso, seconds to wait for the browser login  +to complete +--usernameUSERNAMEThe username to use for authentication + diff --git a/airflow-ctl/docs/start.rst b/airflow-ctl/docs/start.rst index fc879092c8c48..d3f1d96c71307 100644 --- a/airflow-ctl/docs/start.rst +++ b/airflow-ctl/docs/start.rst @@ -44,7 +44,18 @@ OR export AIRFLOW_CLI_TOKEN= airflowctl auth login --api-url --env -This command will create a new airflowctl environment with the specified username and password. +OR, when the auth manager of the API server offers browser login (for example to sign in through SSO): + +.. code-block:: bash + + airflowctl auth login --sso --api-url + +This opens the login page of the API server in your browser and waits for the login to complete. +The browser redirects back to airflowctl on ``127.0.0.1``, so it must run on the same machine. +On a remote machine, pass ``--no-browser`` to only print the login URL, and forward the port +of the printed ``redirect_uri`` from your local machine, for example with ``ssh -L``. + +Each of these commands creates a new airflowctl environment with the given credentials. You can then use the following command to start the airflowctl environment: .. code-block:: bash diff --git a/airflow-ctl/src/airflowctl/api/operations.py b/airflow-ctl/src/airflowctl/api/operations.py index 4e0c9f962ceb9..e41005dd33725 100644 --- a/airflow-ctl/src/airflowctl/api/operations.py +++ b/airflow-ctl/src/airflowctl/api/operations.py @@ -43,6 +43,7 @@ BulkBodyVariableBody, BulkResponse, ClearTaskInstancesBody, + CliLoginMetadataResponse, Config, ConnectionBody, ConnectionCollectionResponse, @@ -282,6 +283,43 @@ def login_with_username_and_password(self, login: LoginBody) -> LoginResponse | self.client.post("/token/cli", json=login.model_dump(mode="json")).content ) + def get_cli_login_metadata(self, api_url: str) -> CliLoginMetadataResponse: + """Return the login methods the API server offers, or none for servers without CLI login discovery.""" + try: + response = self.client.get( + f"{api_url.rstrip('/')}/api/v2/auth/cli/metadata", + extensions={"airflowctl_suppress_error_log": True}, + ) + except httpx.HTTPStatusError as e: + if e.response.status_code != 404: + raise + return CliLoginMetadataResponse( + grant_types_supported=[], + code_challenge_methods_supported=[], + authorization_endpoint=None, + device_authorization_endpoint=None, + token_endpoint=None, + ) + return CliLoginMetadataResponse.model_validate_json(response.content) + + def login_with_authorization_code( + self, token_url: str, code: str, code_verifier: str, redirect_uri: str + ) -> LoginResponse: + """Redeem an authorization code from the browser login for a CLI token.""" + request = self.client.build_request( + "POST", + token_url, + json={ + "grant_type": "authorization_code", + "code": code, + "code_verifier": code_verifier, + "redirect_uri": redirect_uri, + }, + ) + # send() skips the retries of Client.request(): a code is single use, and identity providers can + # treat a replayed code as an attack and revoke the session. + return LoginResponse.model_validate_json(self.client.send(request).content) + # Operations class AssetsOperations(BaseOperations): diff --git a/airflow-ctl/src/airflowctl/ctl/cli_config.py b/airflow-ctl/src/airflowctl/ctl/cli_config.py index 1f5dd5b269100..69e69f4442a20 100755 --- a/airflow-ctl/src/airflowctl/ctl/cli_config.py +++ b/airflow-ctl/src/airflowctl/ctl/cli_config.py @@ -301,6 +301,30 @@ def _load_help_texts_yaml() -> dict[str, dict[str, str]]: dest="password", help="The password to use for authentication", ) +ARG_AUTH_SSO = Arg( + flags=("--sso",), + dest="sso", + default=False, + action="store_true", + help="Log in through the browser with the identity provider of the API server, if it offers it", +) +ARG_AUTH_NO_BROWSER = Arg( + flags=("--no-browser",), + dest="no_browser", + default=False, + action="store_true", + help=( + "With --sso, print the login URL instead of opening a browser. The browser must still reach " + "this machine on 127.0.0.1, for example through 'ssh -L' when logged in remotely" + ), +) +ARG_AUTH_TIMEOUT = Arg( + flags=("--timeout",), + type=float, + default=300, + dest="timeout", + help="With --sso, seconds to wait for the browser login to complete", +) # Dag Commands Args ARG_DAG_ID = Arg( @@ -1154,7 +1178,8 @@ def merge_commands( description=( "Log in to the Airflow API server and store the credentials for the selected environment. " "Pass --api-token (or set AIRFLOW_CLI_TOKEN), or --username and --password; " - "missing credentials are prompted for when running in a terminal." + "missing credentials are prompted for when running in a terminal. " + "Pass --sso to log in through the browser when the API server offers it." ), func=lazy_load_command("airflowctl.ctl.commands.auth_command.login"), args=( @@ -1163,6 +1188,9 @@ def merge_commands( ARG_AUTH_USERNAME, ARG_AUTH_PASSWORD, ARG_AUTH_SKIP_KEYRING, + ARG_AUTH_SSO, + ARG_AUTH_NO_BROWSER, + ARG_AUTH_TIMEOUT, ), ), ActionCommand( @@ -1184,6 +1212,9 @@ def merge_commands( ARG_AUTH_URL, ARG_AUTH_USERNAME, ARG_AUTH_PASSWORD, + ARG_AUTH_SSO, + ARG_AUTH_NO_BROWSER, + ARG_AUTH_TIMEOUT, ), ), ) diff --git a/airflow-ctl/src/airflowctl/ctl/commands/auth_command.py b/airflow-ctl/src/airflowctl/ctl/commands/auth_command.py index cf521cbe7eea0..f79d257d59c7b 100644 --- a/airflow-ctl/src/airflowctl/ctl/commands/auth_command.py +++ b/airflow-ctl/src/airflowctl/ctl/commands/auth_command.py @@ -27,10 +27,40 @@ import keyring import rich from keyring.errors import NoKeyringError +from rich.markup import escape -from airflowctl.api.client import NEW_API_CLIENT, ClientKind, Credentials, provide_api_client +from airflowctl.api.client import NEW_API_CLIENT, Client, ClientKind, Credentials, provide_api_client from airflowctl.api.datamodels.auth_generated import LoginBody from airflowctl.ctl.console_formatting import AirflowConsole +from airflowctl.ctl.utils.browser_login import get_authorization_code, resolve_endpoint +from airflowctl.exceptions import AirflowCtlBrowserLoginException + + +def _get_token_with_browser(args, api_client: Client) -> str: + """Log in through the browser and return the CLI token, if the API server offers it.""" + api_client.refresh_base_url(base_url=args.api_url, kind=ClientKind.AUTH) + metadata = api_client.login.get_cli_login_metadata(args.api_url) + if ( + "authorization_code" not in metadata.grant_types_supported + or "S256" not in metadata.code_challenge_methods_supported + or not metadata.authorization_endpoint + or not metadata.token_endpoint + ): + raise AirflowCtlBrowserLoginException(f"{args.api_url} does not offer browser login (--sso).") + + authorization_url = resolve_endpoint(args.api_url, metadata.authorization_endpoint) + token_url = resolve_endpoint(args.api_url, metadata.token_endpoint) + authorization_code = get_authorization_code( + authorization_url=authorization_url, + open_browser=not args.no_browser, + timeout=args.timeout, + ) + return api_client.login.login_with_authorization_code( + token_url, + code=authorization_code.code, + code_verifier=authorization_code.code_verifier, + redirect_uri=authorization_code.redirect_uri, + ).access_token @provide_api_client(kind=ClientKind.AUTH) @@ -42,6 +72,26 @@ def login(args, api_client=NEW_API_CLIENT) -> None: password = args.password token = args.api_token or os.environ.get("AIRFLOW_CLI_TOKEN") + if args.sso: + if args.api_token or username or password or args.skip_keyring: + rich.print( + "[red]--sso is not compatible with --api-token, --username, --password or --skip-keyring." + ) + sys.exit(1) + try: + api_token = _get_token_with_browser(args, api_client) + except Exception as e: + rich.print(f"[red]Login failed: {escape(str(e))}") + sys.exit(1) + Credentials( + api_url=args.api_url, + api_token=api_token, + api_environment=args.env, + client_kind=ClientKind.AUTH, + ).save() + rich.print(success_message) + return + # If credentials are incomplete, prompt interactively on a real terminal if not token and not (username and password): if not sys.stdin.isatty(): @@ -107,6 +157,17 @@ def login(args, api_client=NEW_API_CLIENT) -> None: @provide_api_client(kind=ClientKind.AUTH) def get_token(args, api_client=NEW_API_CLIENT) -> None: """Generate and print a JWT token for the given credentials to stdout.""" + if args.sso: + if args.username or args.password: + rich.print("[red]--sso is not compatible with --username or --password.") + sys.exit(1) + try: + print(_get_token_with_browser(args, api_client)) + except Exception as e: + rich.print(f"[red]Token generation failed: {escape(str(e))}[/red]") + sys.exit(1) + return + username = args.username or input("Username: ") password = args.password or getpass.getpass("Password: ") diff --git a/airflow-ctl/src/airflowctl/ctl/utils/browser_login.py b/airflow-ctl/src/airflowctl/ctl/utils/browser_login.py new file mode 100644 index 0000000000000..292d03398be9b --- /dev/null +++ b/airflow-ctl/src/airflowctl/ctl/utils/browser_login.py @@ -0,0 +1,157 @@ +# +# 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. +"""Browser login with the authorization code flow and PKCE (RFC 7636), over a loopback redirect (RFC 8252).""" + +from __future__ import annotations + +import base64 +import hashlib +import secrets +import time +import webbrowser +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, HTTPServer +from urllib.parse import parse_qs, urlencode, urljoin, urlsplit + +from rich.console import Console +from rich.markup import escape + +from airflowctl.exceptions import AirflowCtlBrowserLoginException + +CALLBACK_PATH = "/callback" +_CALLBACK_PAGE = ( + b"airflowctl" + b"

Login complete. You can close this window and return to the terminal.

" +) +_CALLBACK_FAILED_PAGE = ( + b"airflowctl

Login failed. Return to the terminal for details.

" +) + + +@dataclass +class AuthorizationCode: + """Authorization code returned to the loopback redirect, with what the token endpoint needs to redeem it.""" + + code: str + code_verifier: str + redirect_uri: str + + +def build_pkce_pair() -> tuple[str, str]: + """Return a PKCE ``(code_verifier, code_challenge)`` pair for the ``S256`` method.""" + code_verifier = secrets.token_urlsafe(64) + digest = hashlib.sha256(code_verifier.encode("ascii")).digest() + code_challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") + return code_verifier, code_challenge + + +def resolve_endpoint(api_url: str, endpoint: str) -> str: + """Return the URL of a login metadata endpoint, which must be on the API server.""" + url = urljoin(api_url, endpoint) + if urlsplit(url)[:2] != urlsplit(api_url)[:2]: + raise AirflowCtlBrowserLoginException(f"{endpoint} is not an endpoint of {api_url}.") + return url + + +def _find_login_error(params: dict[str, str], state: str) -> str | None: + """Return why the callback parameters do not complete this login, or ``None`` if they do.""" + # compare_digest raises TypeError on non-ASCII str. + if not secrets.compare_digest(params.get("state", "").encode(), state.encode()): + return "The login response does not match this login attempt." + if "error" in params: + return f"Login refused: {params.get('error_description') or params['error']}" + if not params.get("code"): + return "The login response contains no authorization code." + return None + + +class _CallbackServer(HTTPServer): + state: str + callback_params: dict[str, str] | None = None + login_error: str | None = None + + +class _CallbackHandler(BaseHTTPRequestHandler): + server: _CallbackServer + # server.timeout only bounds accept(), so an idle connection such as a browser preconnect would block. + timeout = 2 + + def do_GET(self) -> None: + url = urlsplit(self.path) + if url.path != CALLBACK_PATH or self.server.callback_params is not None: + self.send_error(404) + return + params = {key: values[0] for key, values in parse_qs(url.query).items()} + login_error = _find_login_error(params, self.server.state) + self.server.callback_params = params + self.server.login_error = login_error + self.send_response(400 if login_error else 200) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.end_headers() + self.wfile.write(_CALLBACK_FAILED_PAGE if login_error else _CALLBACK_PAGE) + + def log_message(self, format, *args) -> None: + # Request lines contain the authorization code. + pass + + +def get_authorization_code( + *, authorization_url: str, open_browser: bool, timeout: float +) -> AuthorizationCode: + """ + Run the browser part of the login and return the authorization code it produced. + + :param authorization_url: URL of the ``authorization_endpoint`` from the CLI login metadata + :param open_browser: open the login page in the default browser, otherwise only print its URL + :param timeout: seconds to wait for the user to finish logging in + """ + code_verifier, code_challenge = build_pkce_pair() + state = secrets.token_urlsafe(32) + + with _CallbackServer(("127.0.0.1", 0), _CallbackHandler) as server: + server.state = state + redirect_uri = f"http://127.0.0.1:{server.server_address[1]}{CALLBACK_PATH}" + query = urlencode( + { + "redirect_uri": redirect_uri, + "state": state, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } + ) + login_url = f"{authorization_url}?{query}" + # `auth token --sso` prints the token on stdout. + Console(stderr=True, soft_wrap=True).print( + f"[blue]Log in from your browser:[/blue] {escape(login_url)}" + ) + if open_browser: + webbrowser.open(login_url) + + deadline = time.monotonic() + timeout + while server.callback_params is None: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise AirflowCtlBrowserLoginException(f"No login completed within {timeout:g} seconds.") + server.timeout = remaining + server.handle_request() + params = server.callback_params + login_error = server.login_error + + if login_error: + raise AirflowCtlBrowserLoginException(login_error) + return AuthorizationCode(code=params["code"], code_verifier=code_verifier, redirect_uri=redirect_uri) diff --git a/airflow-ctl/src/airflowctl/exceptions.py b/airflow-ctl/src/airflowctl/exceptions.py index 0af8b9fa593cc..835167f61815b 100644 --- a/airflow-ctl/src/airflowctl/exceptions.py +++ b/airflow-ctl/src/airflowctl/exceptions.py @@ -44,3 +44,7 @@ class AirflowCtlConnectionException(AirflowCtlException): class AirflowCtlKeyringException(AirflowCtlException): """Raise when a keyring error occurs while performing an operation.""" + + +class AirflowCtlBrowserLoginException(AirflowCtlException): + """Raise when the browser login of ``auth login --sso`` cannot complete.""" diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py index 6f93d4f20b35c..01780aa555257 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -116,7 +116,7 @@ XComResponse, XComResponseNative, ) -from airflowctl.api.operations import BaseOperations, _build_query_params +from airflowctl.api.operations import BaseOperations, ServerResponseError, _build_query_params from airflowctl.exceptions import AirflowCtlConnectionException if TYPE_CHECKING: @@ -2414,6 +2414,92 @@ def handle_request(request: httpx.Request) -> httpx.Response: ) assert response.access_token == "NO_TOKEN" + def test_get_cli_login_metadata(self): + metadata = { + "grant_types_supported": ["password", "authorization_code"], + "code_challenge_methods_supported": ["S256"], + "authorization_endpoint": "/auth/login/cli", + "device_authorization_endpoint": None, + "token_endpoint": "/auth/token/cli", + } + + def handle_request(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/prefix/api/v2/auth/cli/metadata" + return httpx.Response(200, json=metadata) + + client = make_api_client(transport=httpx.MockTransport(handle_request), kind=ClientKind.AUTH) + + response = client.login.get_cli_login_metadata("http://server/prefix/") + + assert response.model_dump() == metadata + + @pytest.mark.parametrize( + "response", + [ + pytest.param(httpx.Response(404, json={"detail": "Not Found"}), id="json"), + pytest.param(httpx.Response(404, text="Not Found"), id="text"), + ], + ) + def test_get_cli_login_metadata_without_discovery(self, response): + client = make_api_client( + transport=httpx.MockTransport(lambda request: response), kind=ClientKind.AUTH + ) + + metadata = client.login.get_cli_login_metadata("http://server") + + assert metadata.grant_types_supported == [] + assert metadata.token_endpoint is None + + def test_get_cli_login_metadata_server_error(self): + client = make_api_client( + transport=httpx.MockTransport(lambda request: httpx.Response(403, json={"detail": "Forbidden"})), + kind=ClientKind.AUTH, + ) + + with pytest.raises(ServerResponseError): + client.login.get_cli_login_metadata("http://server") + + def test_login_with_authorization_code(self): + def handle_request(request: httpx.Request) -> httpx.Response: + assert str(request.url) == "http://server/auth/token/cli" + assert json.loads(request.content) == { + "grant_type": "authorization_code", + "code": "code", + "code_verifier": "verifier", + "redirect_uri": "http://127.0.0.1:1234/callback", + } + return httpx.Response(201, json={"access_token": "TOKEN"}) + + client = make_api_client(transport=httpx.MockTransport(handle_request), kind=ClientKind.AUTH) + + response = client.login.login_with_authorization_code( + "http://server/auth/token/cli", + code="code", + code_verifier="verifier", + redirect_uri="http://127.0.0.1:1234/callback", + ) + + assert response.access_token == "TOKEN" + + def test_login_with_authorization_code_is_not_retried(self): + requests = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(503, json={"detail": "Unavailable"}) + + client = make_api_client(transport=httpx.MockTransport(handle_request), kind=ClientKind.AUTH) + + with pytest.raises(ServerResponseError): + client.login.login_with_authorization_code( + "http://server/auth/token/cli", + code="code", + code_verifier="verifier", + redirect_uri="http://127.0.0.1:1234/callback", + ) + + assert len(requests) == 1 + class TestXComOperations: """Test suite for XCom operations.""" diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_auth_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_auth_command.py index 9f442f40b9016..cc198a1651723 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_auth_command.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_auth_command.py @@ -22,12 +22,42 @@ from unittest import mock from unittest.mock import patch +import httpx import pytest -from airflowctl.api.client import ClientKind +from airflowctl.api.client import Client, ClientKind from airflowctl.api.datamodels.auth_generated import LoginResponse from airflowctl.ctl import cli_parser from airflowctl.ctl.commands import auth_command +from airflowctl.ctl.utils.browser_login import AuthorizationCode +from airflowctl.exceptions import AirflowCtlBrowserLoginException + +SSO_METADATA = { + "grant_types_supported": ["password", "refresh_token", "authorization_code"], + "code_challenge_methods_supported": ["S256"], + "authorization_endpoint": "/auth/login/cli", + "device_authorization_endpoint": None, + "token_endpoint": "/auth/token/cli", +} +AUTHORIZATION_CODE = AuthorizationCode( + code="code", code_verifier="verifier", redirect_uri="http://127.0.0.1:1234/callback" +) + + +def _make_sso_api_client(metadata: dict) -> Client: + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == "/api/v2/auth/cli/metadata": + return httpx.Response(200, json=metadata) + assert str(request.url) == "http://localhost:8080/auth/token/cli" + assert json.loads(request.content)["code_verifier"] == "verifier" + return httpx.Response(201, json={"access_token": "SSO_TOKEN"}) + + return Client( + base_url="test://server", + transport=httpx.MockTransport(handle_request), + token="", + kind=ClientKind.AUTH, + ) class TestCliAuthCommands: @@ -255,6 +285,115 @@ def test_login_with_username_and_password_no_keyring_backend(self, mock_keyring, api_client=api_client, ) + @patch("airflowctl.ctl.commands.auth_command.get_authorization_code", autospec=True) + @patch("airflowctl.api.client.keyring") + def test_login_with_sso(self, mock_keyring, mock_get_authorization_code): + mock_get_authorization_code.return_value = AUTHORIZATION_CODE + mock_keyring.get_password.return_value = None + + auth_command.login( + self.parser.parse_args( + [ + "auth", + "login", + "--api-url", + "http://localhost:8080", + "--sso", + "--no-browser", + "--timeout", + "30", + ] + ), + api_client=_make_sso_api_client(SSO_METADATA), + ) + + mock_get_authorization_code.assert_called_once_with( + authorization_url="http://localhost:8080/auth/login/cli", + open_browser=False, + timeout=30, + ) + mock_keyring.set_password.assert_called_once_with("airflowctl", "api_token_production", "SSO_TOKEN") + + @pytest.mark.parametrize( + "metadata", + [ + pytest.param({**SSO_METADATA, "grant_types_supported": ["password"]}, id="no-authorization-code"), + pytest.param({**SSO_METADATA, "code_challenge_methods_supported": ["plain"]}, id="no-s256"), + pytest.param({**SSO_METADATA, "authorization_endpoint": None}, id="no-authorization-endpoint"), + pytest.param( + {**SSO_METADATA, "authorization_endpoint": "//evil.example.com/login"}, + id="foreign-authorization-endpoint", + ), + pytest.param( + {**SSO_METADATA, "token_endpoint": "https://evil.example.com/token"}, + id="foreign-token-endpoint", + ), + ], + ) + @patch("airflowctl.ctl.commands.auth_command.get_authorization_code", autospec=True) + def test_login_with_sso_not_offered(self, mock_get_authorization_code, metadata): + with pytest.raises(SystemExit): + auth_command.login( + self.parser.parse_args(["auth", "login", "--api-url", "http://localhost:8080", "--sso"]), + api_client=_make_sso_api_client(metadata), + ) + + mock_get_authorization_code.assert_not_called() + + @pytest.mark.parametrize( + "extra_args", + [["--username", "user"], ["--password", "pass"], ["--api-token", "token"], ["--skip-keyring"]], + ) + @patch("airflowctl.ctl.commands.auth_command.get_authorization_code", autospec=True) + def test_login_with_sso_incompatible_args(self, mock_get_authorization_code, extra_args): + with pytest.raises(SystemExit): + auth_command.login( + self.parser.parse_args( + ["auth", "login", "--api-url", "http://localhost:8080", "--sso", *extra_args] + ), + api_client=_make_sso_api_client(SSO_METADATA), + ) + + mock_get_authorization_code.assert_not_called() + + @pytest.mark.parametrize("extra_args", [["--username", "user"], ["--password", "pass"]]) + @patch("airflowctl.ctl.commands.auth_command.get_authorization_code", autospec=True) + def test_get_token_with_sso_incompatible_args(self, mock_get_authorization_code, extra_args): + with pytest.raises(SystemExit): + auth_command.get_token( + self.parser.parse_args( + ["auth", "token", "--api-url", "http://localhost:8080", "--sso", *extra_args] + ), + api_client=_make_sso_api_client(SSO_METADATA), + ) + + mock_get_authorization_code.assert_not_called() + + @pytest.mark.parametrize( + ("command", "subcommand"), + [pytest.param(auth_command.login, "login", id="login"), (auth_command.get_token, "token")], + ) + @patch("airflowctl.ctl.commands.auth_command.get_authorization_code", autospec=True) + def test_sso_error_is_not_read_as_markup(self, mock_get_authorization_code, command, subcommand): + mock_get_authorization_code.side_effect = AirflowCtlBrowserLoginException("refused [/oops]") + + with pytest.raises(SystemExit): + command( + self.parser.parse_args(["auth", subcommand, "--api-url", "http://localhost:8080", "--sso"]), + api_client=_make_sso_api_client(SSO_METADATA), + ) + + @patch("airflowctl.ctl.commands.auth_command.get_authorization_code", autospec=True) + def test_get_token_with_sso(self, mock_get_authorization_code, capsys): + mock_get_authorization_code.return_value = AUTHORIZATION_CODE + + auth_command.get_token( + self.parser.parse_args(["auth", "token", "--api-url", "http://localhost:8080", "--sso"]), + api_client=_make_sso_api_client(SSO_METADATA), + ) + + assert capsys.readouterr().out == "SSO_TOKEN\n" + class TestListEnvs: parser = cli_parser.get_parser() diff --git a/airflow-ctl/tests/airflow_ctl/ctl/utils/__init__.py b/airflow-ctl/tests/airflow_ctl/ctl/utils/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-ctl/tests/airflow_ctl/ctl/utils/__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/airflow-ctl/tests/airflow_ctl/ctl/utils/test_browser_login.py b/airflow-ctl/tests/airflow_ctl/ctl/utils/test_browser_login.py new file mode 100644 index 0000000000000..4fefece729490 --- /dev/null +++ b/airflow-ctl/tests/airflow_ctl/ctl/utils/test_browser_login.py @@ -0,0 +1,195 @@ +# +# 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 + +import base64 +import hashlib +import queue +import socket +import threading +import urllib.error +import urllib.request +from unittest.mock import patch +from urllib.parse import parse_qs, urlencode, urlsplit + +import pytest + +from airflowctl.ctl.utils.browser_login import build_pkce_pair, get_authorization_code, resolve_endpoint +from airflowctl.exceptions import AirflowCtlBrowserLoginException + + +def _fake_browser(build_callback_query, pages): + """Return a ``webbrowser.open`` stand-in that calls back and puts ``(status, body)`` on ``pages``.""" + + def open_login_page(login_url): + params = {key: values[0] for key, values in parse_qs(urlsplit(login_url).query).items()} + callback_url = f"{params['redirect_uri']}?{urlencode(build_callback_query(params))}" + + def call_back(): + redirect = urlsplit(params["redirect_uri"]) + with pytest.raises(urllib.error.HTTPError): + urllib.request.urlopen(f"http://{redirect.netloc}/favicon.ico") + try: + response = urllib.request.urlopen(callback_url) + except urllib.error.HTTPError as e: + response = e + pages.put((response.status, response.read())) + + threading.Thread(target=call_back, daemon=True).start() + return True + + return open_login_page + + +def test_resolve_endpoint(): + assert ( + resolve_endpoint("https://airflow.example.com/prefix", "/prefix/auth/token/cli") + == "https://airflow.example.com/prefix/auth/token/cli" + ) + + +def test_build_pkce_pair(): + code_verifier, code_challenge = build_pkce_pair() + + assert 43 <= len(code_verifier) <= 128 + expected = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).rstrip(b"=").decode() + assert code_challenge == expected + + +@patch("airflowctl.ctl.utils.browser_login.webbrowser.open", autospec=True) +def test_get_authorization_code(mock_open, capsys): + seen_params = {} + pages = queue.Queue() + + def build_callback_query(params): + seen_params.update(params) + return {"code": "the-code", "state": params["state"]} + + mock_open.side_effect = _fake_browser(build_callback_query, pages) + + result = get_authorization_code( + authorization_url="https://airflow.example.com/prefix/auth/login/cli", + open_browser=True, + timeout=10, + ) + + printed = capsys.readouterr() + assert printed.out == "" + assert printed.err == f"Log in from your browser: {mock_open.call_args.args[0]}\n" + login_url = urlsplit(mock_open.call_args.args[0]) + assert (login_url.netloc, login_url.path) == ("airflow.example.com", "/prefix/auth/login/cli") + assert seen_params["code_challenge_method"] == "S256" + assert result.code == "the-code" + assert result.redirect_uri == seen_params["redirect_uri"] + assert urlsplit(result.redirect_uri).hostname == "127.0.0.1" + challenge = base64.urlsafe_b64encode(hashlib.sha256(result.code_verifier.encode()).digest()) + assert challenge.rstrip(b"=").decode() == seen_params["code_challenge"] + status, page = pages.get(timeout=5) + assert status == 200 + assert b"Login complete" in page + + +@pytest.mark.parametrize( + ("build_callback_query", "message"), + [ + pytest.param( + lambda params: {"code": "the-code", "state": "forged"}, + "does not match this login attempt", + id="state-mismatch", + ), + pytest.param( + lambda params: {"code": "the-code", "state": "forg\u00e9"}, + "does not match this login attempt", + id="non-ascii-state", + ), + pytest.param( + lambda params: { + "error": "access_denied", + "error_description": "User refused", + "state": params["state"], + }, + "Login refused: User refused", + id="error", + ), + pytest.param( + lambda params: {"state": params["state"]}, + "contains no authorization code", + id="no-code", + ), + ], +) +@patch("airflowctl.ctl.utils.browser_login.webbrowser.open", autospec=True) +def test_get_authorization_code_rejects_callback(mock_open, build_callback_query, message): + pages = queue.Queue() + mock_open.side_effect = _fake_browser(build_callback_query, pages) + + with pytest.raises(AirflowCtlBrowserLoginException, match=message): + get_authorization_code( + authorization_url="http://localhost:8080/auth/login/cli", + open_browser=True, + timeout=10, + ) + + status, page = pages.get(timeout=5) + assert status == 400 + assert b"Login failed" in page + assert b"Login complete" not in page + + +@patch("airflowctl.ctl.utils.browser_login.webbrowser.open", autospec=True) +def test_get_authorization_code_drops_idle_connection(mock_open): + closed_by_server = [] + + def open_login_page(login_url): + params = {key: values[0] for key, values in parse_qs(urlsplit(login_url).query).items()} + redirect = urlsplit(params["redirect_uri"]) + + def call_back(): + with socket.create_connection((redirect.hostname, redirect.port), timeout=10) as idle: + try: + closed_by_server.append(idle.recv(1) == b"") + except TimeoutError: + closed_by_server.append(False) + query = urlencode({"code": "the-code", "state": params["state"]}) + urllib.request.urlopen(f"{params['redirect_uri']}?{query}").read() + + threading.Thread(target=call_back, daemon=True).start() + return True + + mock_open.side_effect = open_login_page + + result = get_authorization_code( + authorization_url="http://localhost:8080/auth/login/cli", + open_browser=True, + timeout=30, + ) + + assert closed_by_server == [True] + assert result.code == "the-code" + + +@patch("airflowctl.ctl.utils.browser_login.webbrowser.open", autospec=True) +def test_get_authorization_code_no_browser_times_out(mock_open): + with pytest.raises(AirflowCtlBrowserLoginException, match="No login completed within 0.2 seconds"): + get_authorization_code( + authorization_url="http://localhost:8080/auth/login/cli", + open_browser=False, + timeout=0.2, + ) + + mock_open.assert_not_called()