baptistegh commented on code in PR #75938: URL: https://github.com/apache/airflow/pull/75938#discussion_r4440928906
########## airflow-ctl/src/airflowctl/ctl/utils/browser_login.py: ########## @@ -0,0 +1,158 @@ +# +# 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" +# Sent before the code is exchanged for a token, so it must not claim that the login worked. +_CALLBACK_PAGE = ( + b"<!doctype html><title>airflowctl</title>" + b"<p>You can close this window and return to the terminal to see the result of the login.</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) Review Comment: Done. A callback without the login's `state` now gets the failure page and the wait continues. A real `state` mismatch now shows up as a timeout instead of the "does not match this login attempt" message. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
