Add browser login to airflowctl

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

airflow-ctl/docs/start.rst

@@ -44,7 +44,18 @@ OR
   export AIRFLOW_CLI_TOKEN=<token>
   airflowctl auth login --api-url <api_url> --env <env_name>
 
-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 <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

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):

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,
         ),
     ),
 )

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: ")
 

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"<!doctype html><title>airflowctl</title>"
+    b"<p>Login complete. You can close this window and return to the terminal.</p>"
+)
+_CALLBACK_FAILED_PAGE = (
+    b"<!doctype html><title>airflowctl</title><p>Login failed. Return to the terminal for details.</p>"
+)
+
+
+@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)

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."""