diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py index 91cb6b5f63..0d75a9b839 100644 --- a/graphistry/arrow_uploader.py +++ b/graphistry/arrow_uploader.py @@ -1,6 +1,6 @@ from typing import List, Optional, Dict, Any -import base64, io, json, pyarrow as pa, requests, sys +import io, json, pyarrow as pa, requests, sys from graphistry.privacy import Mode, Privacy, ModeAction from graphistry.otel import inject_trace_headers @@ -14,48 +14,14 @@ serialize_edge_encodings ) -from .exceptions import TokenExpireException +from .exceptions import TokenExpireException, SsoPendingException from .validate.validate_encodings import validate_encodings -from .utils.requests import log_requests_error +from .utils.requests import log_requests_error, switch_org_request, OrgSwitchError, OrgSwitchIdpChallenge from .util import setup_logger from graphistry.models.types import ValidationParam logger = setup_logger(__name__) -def _personal_org_from_jwt(token: str) -> Optional[str]: - """Extract the username claim from a JWT payload, used as a personal-org slug. - - Trust chain: callers pass a JWT just received from the authenticated - /api/v2/o/sso/oidc/jwt/{state}/ endpoint in the same exchange, so we decode - the payload without local signature verification. The server re-validates - the token signature on every subsequent request — an incorrect username - can at worst route the session to an org the user isn't a member of (server - rejects), not grant unauthorized access. Do NOT reuse this helper for - tokens received from outside that trust chain. - - Server contract: ``personal_org.slug == jwt_payload.username`` for users - auto-provisioned via SSO. - - Returns None on any decode/parse failure or missing/non-string username. - """ - try: - parts = token.split('.') - if len(parts) < 2: - return None - # base64 padding: only the missing chars (0, 2, or 3); never extra. - # Inputs whose stripped length is 1 mod 4 are invalid b64; the decode - # call below will raise and the caller-level except returns None. - segment = parts[1] + '=' * (-len(parts[1]) % 4) - payload = json.loads(base64.urlsafe_b64decode(segment)) - if not isinstance(payload, dict): - return None - username = payload.get('username') - return username if isinstance(username, str) and username else None - except Exception as exc: - logger.debug("@_personal_org_from_jwt: failed to extract username: %s", exc) - return None - - class ArrowUploader: def __init__( @@ -251,42 +217,27 @@ def certificate_validation(self, certificate_validation): self.__certificate_validation = certificate_validation - ########################################################################3 - - # @property - # def sso_state(self) -> str: - # return getattr(self, '__sso_state', "") - - ########################################################################3 - - # @property - # def sso_auth_url(self) -> str: - # return getattr(self, '__sso_auth_url') - - ########################################################################3 - - def _switch_org(self, org_name: Optional[str], token: Optional[str]) -> None: if not org_name or not token: return - last = self._client_session._last_switched_org_token - if last == (org_name, token): + if self._client_session.is_org_verified(token, org_name): return try: - switch_url = f"{self.server_base_path}/api/v2/o/{org_name}/switch/" - response = requests.post( - switch_url, - data={'slug': org_name}, - headers=inject_trace_headers({'Authorization': f'Bearer {token}'}), - verify=self.certificate_validation, - ) - log_requests_error(response) - self._client_session._last_switched_org_token = (org_name, token) - from .pygraphistry import PyGraphistry - if PyGraphistry.session is self._client_session: - PyGraphistry.session._last_switched_org_token = (org_name, token) + switch_org_request(self.server_base_path, org_name, token, self.certificate_validation) + except OrgSwitchError as exc: + logger.warning("%s; not recording switch", exc) + return + except OrgSwitchIdpChallenge as exc: + logger.warning("%s; not recording switch", exc) + return except Exception as exc: logger.warning("Failed to switch organization %s: %s", org_name, exc) + return + + self._client_session.mark_org_verified(token, org_name) + from .pygraphistry import PyGraphistry + if PyGraphistry.session is self._client_session: + PyGraphistry.session.mark_org_verified(token, org_name) def login(self, username, password, org_name=None): @@ -344,8 +295,7 @@ def _finalize_login(self, out: requests.Response, org_name: Optional[str]) -> 'A if not logged_in_org_name: # no active_organization in JWT payload raise Exception("You are not authorized to the organization '{}', or server does not support organization, please omit org_name parameter".format(org_name)) else: - # if JWT response with org_name different than the pass in org_name - # => org_name not found and return default organization (currently is personal org) + # JWT's active org differs from the requested org_name -> treat as not found. if logged_in_org_name != org_name: raise Exception("Login Organization is not found in your organization") @@ -431,6 +381,40 @@ def sso_login(self, org_name: Optional[str] = None, idp_name: Optional[str] = No return self + def _personal_org_from_api(self, token: str) -> Optional[str]: + """Look up the caller's personal org via GET /api/v2/my/organizations/. + + Trust chain: token is the one just returned by the SSO exchange in + this same flow. Replaces the old JWT-username-slugify reconstruction + (unreliable: couldn't replicate server-side slug-collision suffixing + and depended on byte-for-byte parity with Django's slugify), with a + real server-confirmed lookup, filtered on the ``is_personal`` flag. + + Returns None on any request/parse failure or if no personal org is found. + """ + try: + out = requests.get( + f'{self.server_base_path}/api/v2/my/organizations/', + params={'limit': 1000}, + verify=self.certificate_validation, + headers=inject_trace_headers({'Authorization': f'Bearer {token}'}) + ) + if not (200 <= out.status_code < 300): + return None + payload = out.json() + orgs = payload.get('results', payload) if isinstance(payload, dict) else payload + if not isinstance(orgs, list): + return None + for org in orgs: + if isinstance(org, dict) and org.get('is_personal'): + slug = org.get('slug') + if isinstance(slug, str) and slug: + return slug + return None + except Exception as exc: + logger.debug("@_personal_org_from_api: failed to fetch personal org: %s", exc) + return None + def sso_get_token(self, state): """ Koa, 04 May 2022 Use state to get token @@ -450,10 +434,12 @@ def sso_get_token(self, state): json_response = out.json() self.token = None if 'status' not in json_response: - raise Exception(out.text) + # Unrecognized body (proxy error page, DRF detail, ...): retryable, not fatal. + raise SsoPendingException(out.text) if json_response['status'] != 'OK': - raise Exception(json_response.get('message', out.text)) + # Normal "not ready yet" answer while the browser login is in flight. + raise SsoPendingException(json_response.get('message', out.text)) data = json_response.get('data', {}) token_value = data.get('token') @@ -464,11 +450,19 @@ def sso_get_token(self, state): active_org = data.get('active_organization') slug = active_org.get('slug') if isinstance(active_org, dict) else None + # Defense-in-depth: treat the reserved site-wide sentinel org ('SITE') as unbound -- no dataset can ever be created under it. + if isinstance(slug, str) and slug.upper() == 'SITE': + logger.warning( + "SSO returned active_organization=%r, the reserved " + "site-wide org; treating as unbound and falling back to " + "personal org. This usually means an IdP org claim was " + "mis-mapped server-side — verify per-org SSO config.", + slug + ) + slug = None + if slug: - # Layer 1: server-bound active_organization. Caller's intent - # (self.org_name from register(org_name=...) or session) must - # MATCH or be ABSENT. Symmetric with _finalize_login's strict - # check for username/password (line ~309-316). + # Layer 1: server-bound active_organization -- caller's intent (self.org_name) must MATCH or be ABSENT. if self.org_name and self.org_name != slug: raise Exception( f"SSO returned active_organization={slug!r}, but caller " @@ -481,11 +475,7 @@ def sso_get_token(self, state): self.org_name = slug self._switch_org(slug, token_value) elif self.org_name: - # Layer 2: caller-supplied, server silent. Preserve caller - # intent — subsequent authenticated requests will validate org - # membership lazily. WARNING because the asymmetric outcome - # (caller asked, server didn't bind) is operationally - # interesting and worth investigating per-org SSO config. + # Layer 2: caller-supplied, server silent -- preserve caller intent; membership validates lazily on later requests. logger.warning( "SSO did not bind active_organization but caller requested " "org_name=%s; preserving caller value. Subsequent requests " @@ -494,27 +484,26 @@ def sso_get_token(self, state): self.org_name ) else: - # Layer 3: caller didn't ask, server didn't bind. Try - # JWT-derived personal-org fallback for first-login UX. See - # _personal_org_from_jwt for the trust-chain rationale. - fallback = _personal_org_from_jwt(token_value) + # Layer 3: caller didn't ask, server didn't bind -- try the server-confirmed personal-org fallback (see _personal_org_from_api). + fallback = self._personal_org_from_api(token_value) if fallback: logger.info( "SSO did not bind active_organization; falling back to " - "JWT-derived personal org=%s", fallback + "personal org=%s", fallback ) self.org_name = fallback self._switch_org(fallback, token_value) else: # Layer 4: nothing claimed, nothing bound, nothing inferable. logger.info( - "SSO did not bind active_organization and no JWT " - "username present; site-wide SSO login completes with " - "no org binding." + "SSO did not bind active_organization and no personal " + "org found; site-wide SSO login completes with no org " + "binding." ) except Exception as e: - logger.error('Unexpected SSO authentication error: %s', out, exc_info=True) + # debug, not error: fires on every routine "not ready yet" poll, not just genuine failures -- error level would spam the console. + logger.debug('SSO token poll did not complete: %s', out, exc_info=True) raise e return self diff --git a/graphistry/client_session.py b/graphistry/client_session.py index 5b7da8aeac..d44b5d7adb 100644 --- a/graphistry/client_session.py +++ b/graphistry/client_session.py @@ -1,4 +1,6 @@ +import base64 import os +import time from dataclasses import is_dataclass, replace from typing import Any, Optional, Literal, cast, Protocol, TypedDict, Dict, MutableMapping, Type, TypeVar, Union, overload, Iterator, Tuple from functools import lru_cache @@ -33,6 +35,27 @@ ] +def _jwt_exp(token: str) -> Optional[float]: + """Best-effort read of a JWT's "exp" claim, without verifying its signature -- + the token is already trusted (this client received it from the server over + the auth flow it just completed); this only reads a claim out of a token + it already holds, never validates one from an untrusted source.""" + try: + payload_b64 = token.split(".")[1] + padded = payload_b64 + "=" * (-len(payload_b64) % 4) + payload = json.loads(base64.urlsafe_b64decode(padded)) + exp = payload.get("exp") + return float(exp) if exp is not None else None + except Exception: + return None + + +def _jwt_expired(token: str) -> bool: + exp = _jwt_exp(token) + # No readable exp claim -> not expired by this check; let the server reject it if it's actually stale. + return exp is not None and exp <= time.time() + + class ClientSession: """ Holds all configuration and authentication state for a Graphistry client. @@ -79,6 +102,8 @@ def __init__(self) -> None: self.idp_name: Optional[str] = None self.sso_state: Optional[str] = None + # org_name the in-flight SSO login was initiated for: self.org_name still holds the previous org until the switch is confirmed. + self.sso_requested_org_name: Optional[str] = None self.personal_key: Optional[str] = None self.personal_key_id: Optional[str] = None @@ -96,8 +121,37 @@ def __init__(self) -> None: # TODO: Migrate to a pattern like Kusto or Spanner self._bolt_driver: Optional[Any] = None - # Tracks the last (org_slug, jwt_token) we successfully switched to on the Hub server - self._last_switched_org_token: Optional[Tuple[str, str]] = None + # Per-org token cache (org_name -> (token, exp)) so re-switching to any previously verified org skips the server round trip until that token's own exp passes. + self._verified_org_tokens: Dict[str, Tuple[str, Optional[float]]] = {} + + def is_org_verified(self, token: Optional[str], org_name: str) -> bool: + """Has org_name already been switched-to under this exact, still-unexpired token?""" + if not token: + return False + entry = self._verified_org_tokens.get(org_name) + return entry is not None and entry[0] == token and not _jwt_expired(token) + + def mark_org_verified(self, token: Optional[str], org_name: str) -> None: + """Record org_name as switched-to under this token.""" + if not token: + return + self._verified_org_tokens[org_name] = (token, _jwt_exp(token)) + + def get_verified_token(self, org_name: str) -> Optional[str]: + """Return a still-unexpired token previously verified for org_name, if any. + + May differ from the currently-active token (e.g. a later SSO login + for a different org minted a new one) -- callers that use this to + skip a switch should swap it in as the active token. + """ + entry = self._verified_org_tokens.get(org_name) + if entry is None: + return None + token, _exp = entry + if _jwt_expired(token): + del self._verified_org_tokens[org_name] + return None + return token def copy(self) -> "ClientSession": """ diff --git a/graphistry/exceptions.py b/graphistry/exceptions.py index 7e7287c7ca..9697600286 100644 --- a/graphistry/exceptions.py +++ b/graphistry/exceptions.py @@ -18,6 +18,15 @@ class SsoStateInvalidException(SsoException): pass +class SsoPendingException(SsoException): + """ + Raised while the SSO token is not ready yet: the browser login is still in + flight, so the token endpoint answers with a non-OK body status (typically + "State is invalid"). Retryable -- the poll loop keeps waiting on it. + """ + pass + + class TokenExpireException(Exception): """ diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py index 1190217758..bebfedf6cc 100644 --- a/graphistry/pygraphistry.py +++ b/graphistry/pygraphistry.py @@ -1,7 +1,7 @@ from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union, cast, overload from typing_extensions import Literal from graphistry.privacy import Mode, ModeAction -from graphistry.utils.requests import log_requests_error +from graphistry.utils.requests import log_requests_error, switch_org_request from graphistry.plugins_types.hypergraph import HypergraphResult from graphistry.plugins_types.gexf_types import GexfEdgeViz, GexfNodeViz, GexfParseEngine from graphistry.client_session import ( @@ -33,7 +33,7 @@ from . import bolt_util from .plotter import Plotter from .util import in_databricks, setup_logger, in_ipython -from .exceptions import SsoRetrieveTokenTimeoutException, TokenExpireException, SsoStateInvalidException +from .exceptions import SsoRetrieveTokenTimeoutException, TokenExpireException, SsoStateInvalidException, SsoPendingException from .messages import ( MSG_REGISTER_MISSING_PASSWORD, @@ -51,6 +51,10 @@ SSO_GET_TOKEN_ELAPSE_SECONDS = 50 +# Consecutive network-level failures tolerated while polling for the SSO token +# before giving up; a sustained outage should not wait out the whole sso_timeout. +SSO_POLL_MAX_CONSECUTIVE_ERRORS = 3 + _client_mode_enabled = False _is_client_mode_warned = False PyGraphistry: "GraphistryClient" = None # type: ignore[assignment] @@ -233,11 +237,15 @@ def sso_login(self, org_name: Optional[str] = None, idp_name: Optional[str] = No ) self.session._is_authenticated = False + # Track which org this SSO flow was initiated for, distinct from + # self.session.org_name (the still-previous org until this switch + # is confirmed) -- see _sso_get_token(). + self.session.sso_requested_org_name = org_name arrow_uploader = ArrowUploader( client_session=self.session, server_base_path=self.protocol() - + "://" - + self.server(), + + "://" + + self.server(), certificate_validation=self.certificate_validation(), ).sso_login(org_name, idp_name) try: @@ -247,7 +255,11 @@ def sso_login(self, org_name: Optional[str] = None, idp_name: Optional[str] = No raise ValueError("ArrowUploader.sso_login returned no token") self.api_token(token) self.session._is_authenticated = True - self._maybe_switch_org(org_name or self.session.org_name) + resolved_org = org_name or self.session.org_name + if resolved_org: + # Record verification now: a token just SSO-minted for resolved_org must never re-challenge on the switch_org_request() below. + self.session.mark_org_verified(token, resolved_org) + self._maybe_switch_org(resolved_org) arrow_uploader.token = None # type: ignore[assignment] return token @@ -302,9 +314,32 @@ def _handle_auth_url(self, auth_url: str, sso_timeout: Optional[int], sso_opt_in time.sleep(1) elapsed_time = 1 token = None + consecutive_poll_errors = 0 while True: - token, org_name = self._sso_get_token() + try: + token, org_name = self._sso_get_token() + consecutive_poll_errors = 0 + except SsoStateInvalidException: + raise + except SsoPendingException as exc: + # Browser login still in flight; keep waiting until sso_timeout. + logger.debug("SSO token not ready yet: %s", exc) + token = None + org_name = None + except (requests.exceptions.RequestException, ValueError) as exc: + # ValueError covers both json and simplejson decode errors. + # Retry network/parse blips; permanent errors (org mismatch) propagate. + consecutive_poll_errors += 1 + if consecutive_poll_errors > SSO_POLL_MAX_CONSECUTIVE_ERRORS: + logger.error( + "Repeated network errors polling SSO token (%s consecutive): %s", + consecutive_poll_errors, exc + ) + raise + logger.debug("Transient error polling SSO token, will retry: %s", exc) + token = None + org_name = None try: if not token: if elapsed_time % 10 == 1: @@ -329,25 +364,18 @@ def _handle_auth_url(self, auth_url: str, sso_timeout: Optional[int], sso_opt_in self.session.org_name = org_name # finish, set back to None self.session.sso_state = None + self.session.sso_requested_org_name = None print("Successfully logged in") + current_token = self.api_token() + if org_name and current_token: + # Same reasoning as the direct-token branch above: don't re-challenge a token just minted for org_name via SSO. + self.session.mark_org_verified(current_token, org_name) self._maybe_switch_org(org_name) return self.api_token() else: print("Please run graphistry.sso_get_token() to complete the authentication after you have authenticated via SSO") return None else: - # print("Start getting token ...") - # token = None - # for i in range(10): - # token, org_name = self._sso_get_token() - # if token: - # # set org_name to sso org - # self._config['org_name'] = org_name - # print("Successfully logged in") - # return self.api_token() - # print("Keep trying to get token ...") - # time.sleep(5) - print("Please run graphistry.sso_get_token() to complete the authentication") return None @@ -367,12 +395,14 @@ def _sso_get_token(self) -> Tuple[Optional[str], Optional[str]]: raise SsoStateInvalidException("[SSO] Invalid SSO state: NoneType encountered") # print("_sso_get_token : {}".format(state)) + # Pass the SSO-requested org explicitly: self.session.org_name still holds the previous org mid-switch. arrow_uploader = ArrowUploader( client_session=self.session, server_base_path=self.protocol() - + "://" - + self.server(), + + "://" + + self.server(), certificate_validation=self.certificate_validation(), + org_name=self.session.sso_requested_org_name, ).sso_get_token(state) try: @@ -415,7 +445,11 @@ def refresh(self, token: Optional[str] = None, fail_silent: bool = False) -> Opt ) self.api_token(token) self.session._is_authenticated = True - self._maybe_switch_org(self.session.org_name) + refreshed_org = self.session.org_name + if token and refreshed_org and self.session.get_verified_token(refreshed_org): + # Refresh reissues under the same org context; carry the grant to the new token. + self.session.mark_org_verified(token, refreshed_org) + self._maybe_switch_org(refreshed_org) return self.api_token() except Exception as e: @@ -519,18 +553,6 @@ def api_token(self, value: Optional[str] = None) -> Optional[str]: self.session._is_authenticated = False return value - # @staticmethod - # def api_token_refresh_ms(value=None): - # """Set or get the API token refresh interval in milliseconds. - # None and 0 interpreted as no refreshing.""" - - # if value is None: - # return PyGraphistry._config["api_token_refresh_ms"] - - # # setter - # if value is not PyGraphistry._config["api_token_refresh_ms"]: - # PyGraphistry._config["api_token_refresh_ms"] = int(value) - def protocol(self, value: Optional[str] = None) -> str: """Set or get the protocol ('http' or 'https'). Set automatically when using a server alias. @@ -788,8 +810,7 @@ def _maybe_switch_org(self, org_name: Optional[str]) -> None: if not org_name: return token = self.api_token() - last = self.session._last_switched_org_token - if token and last == (org_name, token): + if self.session.is_org_verified(token, org_name): return try: self.switch_org(org_name) @@ -2549,10 +2570,8 @@ def _viz_url(self, info: DatasetInfo, url_params: URLParamsDict) -> str: extra, ) - def _switch_org_url(self, org_name): - hostname = self.session.hostname - protocol = self.session.protocol - return "{}://{}/api/v2/o/{}/switch/".format(protocol, hostname, org_name) + def _base_url(self) -> str: + return "{}://{}".format(self.session.protocol, self.session.hostname) def layout_settings(self, @@ -2697,21 +2716,31 @@ def personal_key_secret(self, value: Optional[str] = None): self.session.personal_key_secret = value.strip() def switch_org(self, value: str): - # print(self._switch_org_url(value)) - response = requests.post( - self._switch_org_url(value), - data={'slug': value}, - headers=inject_trace_headers({'Authorization': f'Bearer {self.api_token()}'}), - verify=self.session.certificate_validation, + org_name = value.strip() + token = self.api_token() + + # Fast path 1: org_name already verified under the current token (a plain switch doesn't reissue the token) -- skip the round trip. + if self.session.is_org_verified(token, org_name): + self.session.org_name = org_name + logger.debug("Already switched to organization %s with current token; skipping request.", org_name) + return + + # Fast path 2: an earlier, still-unexpired token already verified org_name (now superseded as active by a later SSO login elsewhere) -- swap it back in. + cached_token = self.session.get_verified_token(org_name) + if cached_token: + # Direct assign: api_token()'s setter would clear _is_authenticated and force a needless refresh(). + self.session.api_token = cached_token + self.session.org_name = org_name + logger.debug("Reusing previously-verified token for organization %s; skipping request.", org_name) + return + + switch_org_request( + self._base_url(), org_name, token, self.session.certificate_validation # type: ignore[arg-type] ) - log_requests_error(response) - result = self._handle_api_response(response) - if result is True: - self.session.org_name = value.strip() - logger.info("Switched to organization: {}".format(value.strip())) - else: # print the error message - raise Exception(result) + self.session.org_name = org_name + self.session.mark_org_verified(token, org_name) + logger.info("Switched to organization: {}".format(org_name)) def _handle_api_response(self, response): try: diff --git a/graphistry/tests/test_arrow_uploader.py b/graphistry/tests/test_arrow_uploader.py index 867ea66024..c0cd8d9cb9 100644 --- a/graphistry/tests/test_arrow_uploader.py +++ b/graphistry/tests/test_arrow_uploader.py @@ -8,6 +8,7 @@ from graphistry import ArrowUploader from graphistry.pygraphistry import PyGraphistry +from graphistry.exceptions import SsoPendingException # TODO mock requests for testing actual effectful code @@ -472,6 +473,31 @@ def test_sso_login_get_sso_token_ok(self, mock_get): assert au.org_name == 'mock-org' mock_switch.assert_called_once_with('mock-org', '123') + @mock.patch('requests.get') + def test_sso_get_token_state_not_ready_raises_pending(self, mock_get): + # The endpoint answers status != OK on every poll until the browser + # login completes. Must be typed retryable, not a generic failure -- + # the poll loop keys off SsoPendingException to keep waiting. + mock_get.return_value = self._mock_response( + json_data={'status': 'Failed', 'message': 'State is invalid'} + ) + + au = ArrowUploader() + + with self.assertRaises(SsoPendingException) as cm: + au.sso_get_token(state='not-ready-yet') + assert 'State is invalid' in str(cm.exception) + + @mock.patch('requests.get') + def test_sso_get_token_status_less_body_raises_pending(self, mock_get): + # Proxy/DRF error body with no 'status' key: also retryable. + mock_get.return_value = self._mock_response(json_data={'detail': 'Not found.'}) + + au = ArrowUploader() + + with self.assertRaises(SsoPendingException): + au.sso_get_token(state='ignored') + @mock.patch('requests.get') def test_sso_get_token_missing_active_organization_no_caller_org(self, mock_get): # Site-wide SSO: server omits active_organization, caller passed no @@ -672,24 +698,66 @@ def test_sso_get_token_layer1_caller_server_match_proceeds(self, mock_get): mock_switch.assert_called_once_with('mock-org', '123') @mock.patch('requests.get') - def test_sso_get_token_layer3_jwt_username_fallback(self, mock_get): - # Layer 3: caller passed nothing, server didn't bind — JWT-derived - # personal-org slug is used. Pins the first-login-UX path that #1230 - # contributed to the unified design. - - # Construct a JWT with username="newuser". No signature verification - # happens client-side, so the signature segment is just a placeholder. - payload_dict = {'user_id': 42, 'username': 'newuser', 'exp': 9999999999} - payload_b64 = base64.urlsafe_b64encode( - json.dumps(payload_dict).encode() - ).rstrip(b'=').decode() - fake_jwt = f"eyJhbGciOiJIUzI1NiJ9.{payload_b64}.placeholder-signature" + def test_sso_get_token_site_wide_slug_falls_back_to_personal_org(self, mock_get): + # Defense-in-depth for the nexus SSO claim-to-org resolver bug: server + # binds active_organization to the reserved site-wide sentinel org + # (slug 'SITE') instead of a real org or nothing. No dataset can ever + # be created under that org (server-side 403), so the client must NOT + # accept it — treat as unbound and fall back to the server-confirmed + # personal org (via GET /api/v2/my/organizations/), same as the + # Layer-3 no-org-at-all path. + + sso_resp = self._mock_response( + json_data={ + 'status': 'OK', + 'message': 'State is valid', + 'data': { + 'token': '123', + 'active_organization': {'slug': 'SITE', 'is_found': True, 'is_member': True}, + } + }) + orgs_resp = self._mock_response(json_data={ + 'results': [ + {'slug': 'acme-team', 'is_personal': False}, + {'slug': 'site-user-personal', 'is_personal': True}, + ] + }) + + def _route(url, *args, **kwargs): + return orgs_resp if '/api/v2/my/organizations/' in url else sso_resp + mock_get.side_effect = _route + + client = PyGraphistry.client() + client.session.org_name = None + PyGraphistry.session.org_name = None + + au = ArrowUploader(client_session=client.session) + + with mock.patch.object(ArrowUploader, "_switch_org") as mock_switch: + with self.assertLogs('graphistry.arrow_uploader', level='WARNING') as log_ctx: + au.sso_get_token(state='ignored-valid') + + assert au.token == '123' + assert au.org_name == 'site-user-personal' + mock_switch.assert_called_once_with('site-user-personal', '123') + assert any( + record.levelname == 'WARNING' and 'SITE' in record.getMessage() + for record in log_ctx.records + ), "Expected a WARNING log surfacing the rejected site-wide org slug" + + @mock.patch('requests.get') + def test_sso_get_token_site_wide_slug_case_insensitive(self, mock_get): + # Server-side matching is case-insensitive (slug__iexact), so the + # client-side guard must be too. mock_resp = self._mock_response( json_data={ 'status': 'OK', 'message': 'State is valid', - 'data': {'token': fake_jwt}, + 'data': { + 'token': '123', + 'active_organization': {'slug': 'site', 'is_found': True, 'is_member': True}, + } }) mock_get.return_value = mock_resp @@ -702,9 +770,46 @@ def test_sso_get_token_layer3_jwt_username_fallback(self, mock_get): with mock.patch.object(ArrowUploader, "_switch_org") as mock_switch: au.sso_get_token(state='ignored-valid') - assert au.token == fake_jwt + assert au.token == '123' + assert au.org_name is None + mock_switch.assert_not_called() + + @mock.patch('requests.get') + def test_sso_get_token_layer3_personal_org_api_fallback(self, mock_get): + # Layer 3: caller passed nothing, server didn't bind — server-confirmed + # personal-org slug (via GET /api/v2/my/organizations/) is used. + # Replaces the old JWT-username-slugify reconstruction with a real + # lookup filtered on is_personal. + + sso_resp = self._mock_response( + json_data={ + 'status': 'OK', + 'message': 'State is valid', + 'data': {'token': '123'}, + }) + orgs_resp = self._mock_response(json_data={ + 'results': [ + {'slug': 'other-team', 'is_personal': False}, + {'slug': 'newuser', 'is_personal': True}, + ] + }) + + def _route(url, *args, **kwargs): + return orgs_resp if '/api/v2/my/organizations/' in url else sso_resp + mock_get.side_effect = _route + + client = PyGraphistry.client() + client.session.org_name = None + PyGraphistry.session.org_name = None + + au = ArrowUploader(client_session=client.session) + + with mock.patch.object(ArrowUploader, "_switch_org") as mock_switch: + au.sso_get_token(state='ignored-valid') + + assert au.token == '123' assert au.org_name == 'newuser' - mock_switch.assert_called_once_with('newuser', fake_jwt) + mock_switch.assert_called_once_with('newuser', '123') @mock.patch('requests.get') def test_sso_get_token_layer2_takes_precedence_over_layer3(self, mock_get): @@ -737,75 +842,66 @@ def test_sso_get_token_layer2_takes_precedence_over_layer3(self, mock_get): mock_switch.assert_not_called() -def _make_test_jwt(payload_dict: dict) -> str: - payload_b64 = base64.urlsafe_b64encode( - json.dumps(payload_dict).encode() - ).rstrip(b'=').decode() - return f"eyJhbGciOiJIUzI1NiJ9.{payload_b64}.placeholder-signature" - - -class TestPersonalOrgFromJwt(unittest.TestCase): - """Unit tests for _personal_org_from_jwt — exercises decode/parse failure - modes in isolation, separately from the sso_get_token integration tests.""" - - def test_valid_jwt_with_username_returns_username(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - token = _make_test_jwt({'username': 'alice', 'exp': 9999999999}) - assert _personal_org_from_jwt(token) == 'alice' - - def test_jwt_with_no_dot_returns_none(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - assert _personal_org_from_jwt('not-a-jwt') is None - - def test_empty_string_returns_none(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - assert _personal_org_from_jwt('') is None - - def test_jwt_with_only_one_segment_returns_none(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - assert _personal_org_from_jwt('header-only') is None - - def test_malformed_base64_payload_returns_none(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - # Use chars that are not valid base64 (e.g. "!!!") - assert _personal_org_from_jwt('header.!!!invalid_b64!!!.sig') is None - - def test_non_json_payload_returns_none(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - not_json = base64.urlsafe_b64encode(b'not even close to json').rstrip(b'=').decode() - assert _personal_org_from_jwt(f"header.{not_json}.sig") is None - - def test_non_dict_payload_returns_none(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - # JWT with a JSON ARRAY (not object) as payload — uncommon but possible. - arr_payload = base64.urlsafe_b64encode(json.dumps(['not', 'a', 'dict']).encode()).rstrip(b'=').decode() - assert _personal_org_from_jwt(f"header.{arr_payload}.sig") is None - - def test_payload_without_username_field_returns_none(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - token = _make_test_jwt({'user_id': 1, 'email': 'a@b.com', 'exp': 9999999999}) - assert _personal_org_from_jwt(token) is None - - def test_payload_with_non_string_username_returns_none(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - token = _make_test_jwt({'username': 12345, 'exp': 9999999999}) - assert _personal_org_from_jwt(token) is None - - def test_payload_with_empty_string_username_returns_none(self): - from graphistry.arrow_uploader import _personal_org_from_jwt - token = _make_test_jwt({'username': '', 'exp': 9999999999}) - assert _personal_org_from_jwt(token) is None - - def test_payload_length_multiple_of_4_after_rstrip_decodes_correctly(self): - # Pin the corrected b64 padding formula: when stripped len % 4 == 0, - # we must add ZERO pad chars (not 4). Use a payload that produces a - # rstripped b64 of length multiple of 4. - from graphistry.arrow_uploader import _personal_org_from_jwt - # 9-byte body → b64 length 12 → rstrip removes 0 pads → len 12, mod4 0 - payload_dict = {'username': 'u'} # tiny payload - # Pad the payload until rstripped length is a multiple of 4. A 12-char - # rstripped output corresponds to a 9-byte payload. {"username":"u"} is - # 16 bytes — let's just trust the formula and exercise it explicitly. - token = _make_test_jwt(payload_dict) - # The exact byte length will vary; the corrected formula is robust to all. - assert _personal_org_from_jwt(token) == 'u' +class TestPersonalOrgFromApi(unittest.TestCase): + """Unit tests for ArrowUploader._personal_org_from_api — exercises the + GET /api/v2/my/organizations/ lookup in isolation, separately from the + sso_get_token integration tests.""" + + def _mock_response(self, json_data, status_code=200): + mock_resp = mock.Mock() + mock_resp.status_code = status_code + mock_resp.json.return_value = json_data + mock_resp.text = json.dumps(json_data) + return mock_resp + + @mock.patch('requests.get') + def test_returns_slug_of_is_personal_org(self, mock_get): + mock_get.return_value = self._mock_response({ + 'results': [ + {'slug': 'acme-team', 'is_personal': False}, + {'slug': 'alice-personal', 'is_personal': True}, + ] + }) + au = ArrowUploader() + assert au._personal_org_from_api('tok') == 'alice-personal' + + @mock.patch('requests.get') + def test_non_paginated_list_response(self, mock_get): + # Defensive shape guard: tolerate a plain list too (no 'results' key). + mock_get.return_value = self._mock_response([ + {'slug': 'alice-personal', 'is_personal': True}, + ]) + au = ArrowUploader() + assert au._personal_org_from_api('tok') == 'alice-personal' + + @mock.patch('requests.get') + def test_no_personal_org_in_results_returns_none(self, mock_get): + mock_get.return_value = self._mock_response({ + 'results': [{'slug': 'acme-team', 'is_personal': False}] + }) + au = ArrowUploader() + assert au._personal_org_from_api('tok') is None + + @mock.patch('requests.get') + def test_empty_results_returns_none(self, mock_get): + mock_get.return_value = self._mock_response({'results': []}) + au = ArrowUploader() + assert au._personal_org_from_api('tok') is None + + @mock.patch('requests.get') + def test_non_200_status_returns_none(self, mock_get): + mock_get.return_value = self._mock_response({'detail': 'unauthorized'}, status_code=401) + au = ArrowUploader() + assert au._personal_org_from_api('tok') is None + + @mock.patch('requests.get') + def test_unexpected_shape_returns_none(self, mock_get): + mock_get.return_value = self._mock_response({'results': 'not-a-list'}) + au = ArrowUploader() + assert au._personal_org_from_api('tok') is None + + @mock.patch('requests.get') + def test_request_exception_returns_none(self, mock_get): + mock_get.side_effect = Exception('connection error') + au = ArrowUploader() + assert au._personal_org_from_api('tok') is None diff --git a/graphistry/tests/test_client_session.py b/graphistry/tests/test_client_session.py index 40e42c4327..28725d0664 100644 --- a/graphistry/tests/test_client_session.py +++ b/graphistry/tests/test_client_session.py @@ -1,3 +1,7 @@ +import base64 +import json +import time + import pytest import graphistry from graphistry.client_session import ClientSession @@ -7,6 +11,52 @@ from unittest import mock +def _fake_jwt(exp=None) -> str: + header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode() + payload = {} if exp is None else {"exp": exp} + body = base64.urlsafe_b64encode(json.dumps(payload).encode()).rstrip(b"=").decode() + return f"{header}.{body}.sig" + + +class TestVerifiedOrgTokenCache: + """Per-org, exp-aware token cache backing switch_org()'s client-only fast path.""" + + def test_multiple_orgs_survive_under_distinct_tokens(self): + session = ClientSession() + token_a = _fake_jwt(exp=time.time() + 3600) + token_b = _fake_jwt(exp=time.time() + 3600) + + session.mark_org_verified(token_a, "org-a") + session.mark_org_verified(token_b, "org-b") + + # Both remain retrievable even though token_b is "more current" -- + # a fresh SSO login for org-b doesn't evict org-a's own grant. + assert session.get_verified_token("org-a") == token_a + assert session.get_verified_token("org-b") == token_b + + def test_expired_entry_is_not_reused_and_is_evicted(self): + session = ClientSession() + expired = _fake_jwt(exp=time.time() - 1) + session.mark_org_verified(expired, "org-a") + + assert session.get_verified_token("org-a") is None + # Second call confirms it was actually evicted, not just filtered. + assert "org-a" not in session._verified_org_tokens + + def test_is_org_verified_requires_exact_current_token(self): + session = ClientSession() + token = _fake_jwt(exp=time.time() + 3600) + session.mark_org_verified(token, "org-a") + + assert session.is_org_verified(token, "org-a") is True + assert session.is_org_verified("some-other-token", "org-a") is False + + def test_unknown_org_returns_none(self): + session = ClientSession() + assert session.get_verified_token("never-switched") is None + assert session.is_org_verified("tok", "never-switched") is False + + class TestClientSession: """ Tests for Graphistry client()/register() interactions using the *current* @@ -153,7 +203,7 @@ def test_client_register_updates_last_switched_cache(self, mock_post, mock_switc client = graphistry.client() def fake_switch(org_name, token): - client.session._last_switched_org_token = (org_name, token) + client.session.mark_org_verified(token, org_name) mock_switch_org.side_effect = fake_switch mock_resp = mock.Mock() @@ -172,11 +222,11 @@ def fake_switch(org_name, token): mock_resp.raise_for_status = mock.Mock() mock_post.return_value = mock_resp - assert client.session._last_switched_org_token is None + assert client.session._verified_org_tokens == {} client.register(api=3, username='u', password='p', org_name='mock-org') - assert client.session._last_switched_org_token == ('mock-org', 'tok123') + assert client.session.is_org_verified('tok123', 'mock-org') # --------------------------------------------------------------------- # # Persistence of arbitrary config # diff --git a/graphistry/tests/test_pygraphistry.py b/graphistry/tests/test_pygraphistry.py index c71f5b54aa..fcc8ca0e58 100644 --- a/graphistry/tests/test_pygraphistry.py +++ b/graphistry/tests/test_pygraphistry.py @@ -1,5 +1,9 @@ # -*- coding: utf-8 -*- +import base64 +import json +import time +from typing import Optional import unittest, pytest try: from mock import patch # type: ignore @@ -7,7 +11,8 @@ from unittest.mock import patch import graphistry -from graphistry.pygraphistry import PyGraphistry, GraphistryClient +from graphistry.pygraphistry import PyGraphistry, GraphistryClient, SSO_POLL_MAX_CONSECUTIVE_ERRORS +from graphistry.exceptions import SsoPendingException from graphistry.messages import ( MSG_REGISTER_MISSING_PASSWORD, MSG_REGISTER_MISSING_USERNAME, @@ -86,7 +91,7 @@ def test_refresh_skips_switch_when_cached(mock_refresh): client = graphistry.client() client.session.org_name = "mock-org" client.api_token("tok123") - client.session._last_switched_org_token = ("mock-org", "tok123") + client.session.mark_org_verified("tok123", "mock-org") with patch.object(client, "switch_org") as mock_switch: client.refresh() @@ -103,7 +108,7 @@ def test_refresh_switches_when_org_changes(mock_refresh): client = graphistry.client() client.session.org_name = "new-org" client.api_token("tok123") - client.session._last_switched_org_token = ("old-org", "tok123") + client.session.mark_org_verified("tok123", "old-org") with patch.object(client, "switch_org") as mock_switch: client.refresh() @@ -114,7 +119,7 @@ def test_refresh_switches_when_org_changes(mock_refresh): def test_maybe_switch_org_cached_pair_skips(): client = graphistry.client() client.api_token("tok123") - client.session._last_switched_org_token = ("mock-org", "tok123") + client.session.mark_org_verified("tok123", "mock-org") with patch.object(client, "switch_org") as mock_switch: client._maybe_switch_org("mock-org") @@ -125,7 +130,7 @@ def test_maybe_switch_org_cached_pair_skips(): def test_maybe_switch_org_new_token_switches(): client = graphistry.client() client.api_token("tok123") - client.session._last_switched_org_token = ("mock-org", "old-token") + client.session.mark_org_verified("old-token", "mock-org") with patch.object(client, "switch_org") as mock_switch: client._maybe_switch_org("mock-org") @@ -136,7 +141,7 @@ def test_maybe_switch_org_new_token_switches(): def test_maybe_switch_org_new_org_switches(): client = graphistry.client() client.api_token("tok123") - client.session._last_switched_org_token = ("other-org", "tok123") + client.session.mark_org_verified("tok123", "other-org") with patch.object(client, "switch_org") as mock_switch: client._maybe_switch_org("mock-org") @@ -144,6 +149,178 @@ def test_maybe_switch_org_new_org_switches(): mock_switch.assert_called_once_with("mock-org") +def _fake_jwt(exp: Optional[float] = None) -> str: + """Minimal unsigned-looking JWT with an optional exp claim, for exercising + ClientSession's exp-aware verified-org cache without a real server.""" + header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode() + payload = {} if exp is None else {"exp": exp} + body = base64.urlsafe_b64encode(json.dumps(payload).encode()).rstrip(b"=").decode() + return f"{header}.{body}.sig" + + +def test_switch_org_reuses_earlier_verified_token_for_different_org(): + """Org A was SSO-verified under token1; a later SSO login to org B mints + token2 (now active). Switching back to org A should swap token1 back in + locally instead of hitting the server, since token1 hasn't expired.""" + client = graphistry.client() + token1 = _fake_jwt(exp=time.time() + 3600) + token2 = _fake_jwt(exp=time.time() + 3600) + + client.api_token(token1) + client.session.mark_org_verified(token1, "org-a") + + client.api_token(token2) + client.session.mark_org_verified(token2, "org-b") + client.session._is_authenticated = True + + with patch("graphistry.pygraphistry.switch_org_request") as mock_req: + client.switch_org("org-a") + + mock_req.assert_not_called() + assert client.api_token() == token1 + assert client.session.org_name == "org-a" + # Swapping in an already-verified token must not de-authenticate the session, + # else the next authenticate() refreshes and undoes the skipped round trip. + assert client.session._is_authenticated is True + + +def test_switch_org_does_not_reuse_expired_cached_token(): + client = graphistry.client() + expired_token = _fake_jwt(exp=time.time() - 10) + active_token = _fake_jwt(exp=time.time() + 3600) + + client.api_token(expired_token) + client.session.mark_org_verified(expired_token, "org-a") + client.api_token(active_token) + + with patch("graphistry.pygraphistry.switch_org_request") as mock_req: + client.switch_org("org-a") + + mock_req.assert_called_once() + assert client.api_token() == active_token + + +class TestSsoPollErrorHandling: + """_handle_auth_url's polling loop: keep waiting on SsoPendingException and + network blips, propagate permanent errors.""" + + def _client(self): + client = graphistry.client() + client.session.sso_state = "state-123" + return client + + def test_pending_state_keeps_polling_until_login_completes(self): + # The token endpoint answers "State is invalid" on every poll until the + # user finishes the browser login -- that is the normal wait, not a failure. + outcomes = [ + SsoPendingException("State is invalid"), + SsoPendingException("State is invalid"), + SsoPendingException("State is invalid"), + SsoPendingException("State is invalid"), + ("tok-final", "org-a"), + ] + + def _poll(): + outcome = outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + client = self._client() + with patch.object(client, "_sso_get_token", side_effect=_poll) as mock_get: + with patch("graphistry.pygraphistry.time.sleep"): + with patch.object(client, "api_token", return_value="tok-final"): + with patch.object(client, "_maybe_switch_org"): + out = client._handle_auth_url("http://auth", 30, None) + + # More pending polls than the network-error budget: pending must not be capped. + assert mock_get.call_count == 5 + assert out == "tok-final" + + def test_pending_state_stops_at_timeout(self): + client = self._client() + with patch.object( + client, "_sso_get_token", + side_effect=SsoPendingException("State is invalid") + ): + with patch("graphistry.pygraphistry.time.sleep"): + out = client._handle_auth_url("http://auth", 3, None) + + assert out is None + + def test_status_less_body_is_treated_as_pending(self): + # A proxy/DRF error body with no 'status' key must not abort the login. + outcomes = [ + SsoPendingException('{"detail": "Not found."}'), + ("tok-final", "org-a"), + ] + + def _poll(): + outcome = outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + client = self._client() + with patch.object(client, "_sso_get_token", side_effect=_poll): + with patch("graphistry.pygraphistry.time.sleep"): + with patch.object(client, "api_token", return_value="tok-final"): + with patch.object(client, "_maybe_switch_org"): + out = client._handle_auth_url("http://auth", 30, None) + + assert out == "tok-final" + + def test_permanent_error_propagates_instead_of_timing_out(self): + client = self._client() + boom = Exception("SSO returned active_organization='org-b', but caller requested org_name='org-a'") + + with patch.object(client, "_sso_get_token", side_effect=boom) as mock_get: + with patch("graphistry.pygraphistry.time.sleep"): + with pytest.raises(Exception) as excinfo: + client._handle_auth_url("http://auth", 30, None) + + assert "active_organization" in str(excinfo.value) + # Failed on the first poll rather than spinning out the full timeout. + assert mock_get.call_count == 1 + + def test_transient_network_error_is_retried_then_succeeds(self): + client = self._client() + import requests as _requests + outcomes = [ + _requests.exceptions.ConnectionError("reset"), + ("tok-final", "org-a"), + ] + + def _poll(): + outcome = outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + with patch.object(client, "_sso_get_token", side_effect=_poll): + with patch("graphistry.pygraphistry.time.sleep"): + with patch.object(client, "api_token", return_value="tok-final"): + with patch.object(client, "_maybe_switch_org"): + out = client._handle_auth_url("http://auth", 30, None) + + assert out == "tok-final" + assert client.session.org_name == "org-a" + + def test_sustained_network_errors_give_up_before_timeout(self): + client = self._client() + import requests as _requests + + with patch.object( + client, "_sso_get_token", + side_effect=_requests.exceptions.ConnectionError("down") + ) as mock_get: + with patch("graphistry.pygraphistry.time.sleep"): + with pytest.raises(_requests.exceptions.ConnectionError): + client._handle_auth_url("http://auth", 300, None) + + assert mock_get.call_count == SSO_POLL_MAX_CONSECUTIVE_ERRORS + 1 + + class FakeRequestResponse(object): def __init__(self, response, status_code: int): self.response = response diff --git a/graphistry/tests/utils/test_requests.py b/graphistry/tests/utils/test_requests.py new file mode 100644 index 0000000000..a9f0fa2d00 --- /dev/null +++ b/graphistry/tests/utils/test_requests.py @@ -0,0 +1,141 @@ +# -*- coding: utf-8 -*- +import json +import unittest + +try: + import mock # type: ignore +except ImportError: # pragma: no cover - fallback for stdlib-only envs + from unittest import mock + +from graphistry.arrow_uploader import ArrowUploader +from graphistry.client_session import ClientSession +from graphistry.utils import requests as switch_org_requests_module +from graphistry.utils.requests import ( + OrgSwitchError, + OrgSwitchIdpChallenge, + switch_org_request, +) + + +class FakeResponse(object): + """Minimal requests.Response stand-in for switch_org_request tests.""" + + def __init__(self, status_code: int, body=None, raise_on_json: bool = False): + self.status_code = status_code + self._body = body + self._raise_on_json = raise_on_json + self.text = json.dumps(body) if body is not None else "" + + def json(self): + if self._raise_on_json: + raise json.JSONDecodeError("no json", "", 0) + return self._body + + +class TestSwitchOrgRequest(unittest.TestCase): + @mock.patch.object(switch_org_requests_module.requests, "post") + def test_success(self, mock_post): + mock_post.return_value = FakeResponse(200, {"status": "OK", "data": []}) + # Should not raise + switch_org_request("https://hub.graphistry.com", "my-org", "tok", True) + args, kwargs = mock_post.call_args + assert args[0] == "https://hub.graphistry.com/api/v2/o/my-org/switch/" + assert kwargs["data"] == {"slug": "my-org"} + assert kwargs["headers"]["Authorization"] == "Bearer tok" + assert kwargs["verify"] is True + + @mock.patch.object(switch_org_requests_module.requests, "post") + def test_http_error_status(self, mock_post): + mock_post.return_value = FakeResponse(404, {"status": "Failed", "message": "not found"}) + with self.assertRaises(OrgSwitchError) as cm: + switch_org_request("https://hub.graphistry.com", "no-such-org", "tok", True) + assert cm.exception.status_code == 404 + assert cm.exception.org_name == "no-such-org" + assert "404" in str(cm.exception) + + @mock.patch.object(switch_org_requests_module.requests, "post") + def test_body_status_not_ok_with_detail(self, mock_post): + mock_post.return_value = FakeResponse(200, {"status": "Failed", "message": "not permitted"}) + with self.assertRaises(OrgSwitchError) as cm: + switch_org_request("https://hub.graphistry.com", "not-permitted-org", "tok", True) + assert cm.exception.detail == "not permitted" + assert "not permitted" in str(cm.exception) + + @mock.patch.object(switch_org_requests_module.requests, "post") + def test_unparseable_json_body_treated_as_failure(self, mock_post): + mock_post.return_value = FakeResponse(200, None, raise_on_json=True) + with self.assertRaises(OrgSwitchError) as cm: + switch_org_request("https://hub.graphistry.com", "my-org", "tok", True) + # falls back to body={} -> status missing -> not 'OK' + assert cm.exception.detail == "" + + @mock.patch.object(switch_org_requests_module.requests, "post") + def test_idp_challenge(self, mock_post): + idp = {"auth_url": "https://sso-idp-host/authorize?state=xxx"} + mock_post.return_value = FakeResponse( + 200, {"status": "OK", "data": {"idp": idp}} + ) + with self.assertRaises(OrgSwitchIdpChallenge) as cm: + switch_org_request("https://hub.graphistry.com", "sso-org", "tok", True) + assert cm.exception.org_name == "sso-org" + assert cm.exception.idp == idp + assert "sso-org" in str(cm.exception) + assert "auth_url" in str(cm.exception) + + +class TestArrowUploaderSwitchOrg(unittest.TestCase): + def _uploader(self): + return ArrowUploader(client_session=ClientSession()) + + def test_noop_without_org_or_token(self): + au = self._uploader() + with mock.patch("graphistry.arrow_uploader.switch_org_request") as mock_req: + au._switch_org(None, "tok") + au._switch_org("org", None) + au._switch_org(None, None) + mock_req.assert_not_called() + + def test_skips_when_already_switched(self): + au = self._uploader() + au._client_session.mark_org_verified("tok", "my-org") + with mock.patch("graphistry.arrow_uploader.switch_org_request") as mock_req: + au._switch_org("my-org", "tok") + mock_req.assert_not_called() + + def test_success_records_last_switched(self): + au = self._uploader() + with mock.patch("graphistry.arrow_uploader.switch_org_request") as mock_req: + au._switch_org("my-org", "tok") + mock_req.assert_called_once_with(au.server_base_path, "my-org", "tok", au.certificate_validation) + assert au._client_session.is_org_verified("tok", "my-org") + + def test_org_switch_error_does_not_record(self): + au = self._uploader() + with mock.patch( + "graphistry.arrow_uploader.switch_org_request", + side_effect=OrgSwitchError("my-org", 404), + ): + au._switch_org("my-org", "tok") + assert au._client_session._verified_org_tokens == {} + + def test_idp_challenge_does_not_record(self): + au = self._uploader() + with mock.patch( + "graphistry.arrow_uploader.switch_org_request", + side_effect=OrgSwitchIdpChallenge("my-org", {"auth_url": "https://x"}), + ): + au._switch_org("my-org", "tok") + assert au._client_session._verified_org_tokens == {} + + def test_generic_exception_does_not_record(self): + au = self._uploader() + with mock.patch( + "graphistry.arrow_uploader.switch_org_request", + side_effect=RuntimeError("boom"), + ): + au._switch_org("my-org", "tok") + assert au._client_session._verified_org_tokens == {} + + +if __name__ == "__main__": + unittest.main() diff --git a/graphistry/utils/requests.py b/graphistry/utils/requests.py index 1ec9912512..8d93766e0d 100644 --- a/graphistry/utils/requests.py +++ b/graphistry/utils/requests.py @@ -1,6 +1,9 @@ import json +from typing import Any, Dict + import requests +from graphistry.otel import inject_trace_headers from graphistry.util import setup_logger logger = setup_logger(__name__) @@ -14,3 +17,73 @@ def log_requests_error(resp: requests.Response) -> None: logger.error("HTTP %s error - response content (JSON): %s", resp.status_code, json.dumps(error_content, indent=2)) except json.JSONDecodeError: logger.error("HTTP %s error - response content (text): %s", resp.status_code, resp.text) + + +class OrgSwitchError(Exception): + """Org switch endpoint responded with a non-2xx status or a non-OK body status.""" + + def __init__(self, org_name: str, status_code: int, detail: str = '') -> None: + self.org_name = org_name + self.status_code = status_code + self.detail = detail + msg = f"Org switch to '{org_name}' failed with HTTP {status_code}" + if detail: + msg += f": {detail}" + super().__init__(msg) + + +class OrgSwitchIdpChallenge(Exception): + """Server-side quirk: for an org the caller is a plain member of (not owner/admin) + that has its own SSO IDP configured, the switch endpoint still returns + status=='OK' but does NOT actually switch -- it instead returns an SSO + challenge (data['idp']) requiring a fresh org-scoped SSO login before the + switch is honored. Treating that as success would silently leave the + caller on their previous org. + """ + + def __init__(self, org_name: str, idp: Dict[str, Any]) -> None: # hygiene-ok: explicit-any -- idp is arbitrary server-provided JSON + self.org_name = org_name + self.idp = idp + super().__init__( + "Switching to organization '{}' requires SSO re-authentication " + "(idp: {}) -- the server did not actually switch. Complete the " + "auth_url in the response's data['idp'], or use " + "sso_login(org_name='{}') instead of switch_org().".format( + org_name, list(idp.keys()), org_name + ) + ) + + +def switch_org_request(base_url: str, org_name: str, token: str, verify: bool) -> None: + """POST to the org-switch endpoint and validate the response. + + Single source of truth for both GraphistryClient.switch_org and + ArrowUploader._switch_org: same URL shape, same HTTP/body status check, + same SSO-idp-challenge detection. Raises OrgSwitchError or + OrgSwitchIdpChallenge on any failure; returns None on a confirmed switch. + Callers decide their own success/failure side effects (raising vs. + logging-and-swallowing, which session state to update). + """ + switch_url = f"{base_url}/api/v2/o/{org_name}/switch/" + response = requests.post( + switch_url, + data={'slug': org_name}, + headers=inject_trace_headers({'Authorization': f'Bearer {token}'}), + verify=verify, + ) + log_requests_error(response) + + if not (200 <= response.status_code < 300): + raise OrgSwitchError(org_name, response.status_code) + + try: + body = response.json() + except Exception: + body = {} + + if body.get('status') != 'OK': + raise OrgSwitchError(org_name, response.status_code, detail=body.get('message', '')) + + data = body.get('data', {}) + if isinstance(data, dict) and data.get('idp'): + raise OrgSwitchIdpChallenge(org_name, data['idp'])