diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py index 91cb6b5f63..22f7774792 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 @@ -22,40 +22,6 @@ 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__( @@ -431,6 +397,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 @@ -464,6 +464,22 @@ 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: nexus's SSO claim-to-org resolver can (bug) + # bind a user to the reserved site-wide sentinel org (slug + # 'SITE') when an IdP org claim happens to case-insensitively + # match it. No dataset can ever be created under that org (server + # rejects with 403), so treat it the same as "no org bound" and + # fall through to the JWT-derived personal-org fallback below. + 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 @@ -495,22 +511,22 @@ def sso_get_token(self, state): ) 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) + # server-confirmed personal-org fallback for first-login UX. + # See _personal_org_from_api for the trust-chain rationale. + 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: diff --git a/graphistry/tests/test_arrow_uploader.py b/graphistry/tests/test_arrow_uploader.py index 867ea66024..b975e9b6d8 100644 --- a/graphistry/tests/test_arrow_uploader.py +++ b/graphistry/tests/test_arrow_uploader.py @@ -672,24 +672,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 +744,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 +816,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