diff --git a/pyatlan/client/aio/sso.py b/pyatlan/client/aio/sso.py index 3a39c2325..32454324e 100644 --- a/pyatlan/client/aio/sso.py +++ b/pyatlan/client/aio/sso.py @@ -3,7 +3,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, List +from typing import List, cast from pydantic.v1 import validate_arguments @@ -13,15 +13,15 @@ SSOCreateGroupMapping, SSODeleteGroupMapping, SSOGetAllGroupMappings, + SSOGetAllIdentityProviders, SSOGetGroupMapping, SSOUpdateGroupMapping, + SSOUpdateIdentityProvider, + normalize_signing_certificate, ) from pyatlan.errors import ErrorCode from pyatlan.model.group import AtlanGroup -from pyatlan.model.sso import SSOMapper - -if TYPE_CHECKING: - pass +from pyatlan.model.sso import SSOMapper, SSOProvider class AsyncSSOClient: @@ -142,3 +142,81 @@ async def delete_group_mapping(self, sso_alias: str, group_map_id: str) -> None: ) raw_json = await self._client._call_api(endpoint, request_obj=request_obj) return raw_json + + async def get_all_identity_providers(self) -> List[SSOProvider]: + """ + Retrieves all SSO identity providers configured on the tenant. + + Requires an API token with the admin role or workspace-admin + subrole (READ_TENANT_IDP); other tokens receive a 403. + + :raises AtlanError: on any error during API invocation. + :returns: list of the tenant's SSO identity providers + """ + endpoint, request_obj = SSOGetAllIdentityProviders.prepare_request() + raw_json = await self._client._call_api(endpoint, request_obj=request_obj) + return SSOGetAllIdentityProviders.process_response(raw_json) + + @validate_arguments + async def get_identity_provider(self, sso_alias: str) -> SSOProvider: + """ + Retrieves the SSO identity provider with the given alias. + + :param sso_alias: alias of the SSO provider (e.g. `azure`, `okta`) + :raises AtlanError: on any error during API invocation. + :raises NotFoundError: if no identity provider exists with the given alias. + :returns: the identity provider configuration + """ + for provider in await self.get_all_identity_providers(): + if provider.alias == sso_alias: + return provider + raise ErrorCode.IDP_NOT_FOUND_BY_ALIAS.exception_with_parameters(sso_alias) + + @validate_arguments + async def update_identity_provider(self, provider: SSOProvider) -> SSOProvider: + """ + Updates an SSO identity provider's configuration. + + The backend treats this as a full replacement: always retrieve the + current configuration first (`get_identity_provider()`), modify it, + and pass the complete object here. Sending a partial object may + silently reset fields that were omitted. + + Requires an API token with the admin role or workspace-admin + subrole (UPDATE_TENANT_IDP); other tokens receive a 403. + + :param provider: the complete identity provider configuration to store + :raises AtlanError: on any error during API invocation. + :returns: the identity provider configuration, re-read after the update + """ + endpoint, request_obj = SSOUpdateIdentityProvider.prepare_request(provider) + alias = cast(str, provider.alias) # validated non-empty by prepare_request + await self._client._call_api(endpoint, request_obj=request_obj) + return await self.get_identity_provider(sso_alias=alias) + + @validate_arguments + async def update_signing_certificate( + self, sso_alias: str, certificate: str + ) -> SSOProvider: + """ + Replaces the signing certificate on the given SSO identity + provider, leaving the rest of the configuration untouched. + + API tokens authenticate independently of SSO, so this works even + while SSO logins are failing (for example, after the certificate + expired) - as long as the token already exists. + + :param sso_alias: alias of the SSO provider (e.g. `azure`, `okta`) + :param certificate: new X.509 certificate, as PEM or a single line + of base64; stored as one line with no BEGIN/END lines and no + line breaks + :raises AtlanError: on any error during API invocation. + :raises NotFoundError: if no identity provider exists with the given alias. + :returns: the identity provider configuration, re-read after the update + """ + provider = await self.get_identity_provider(sso_alias=sso_alias) + provider.config = provider.config or {} + provider.config["signingCertificate"] = normalize_signing_certificate( + certificate + ) + return await self.update_identity_provider(provider=provider) diff --git a/pyatlan/client/common/__init__.py b/pyatlan/client/common/__init__.py index 3ba1baa14..6b1ca94a7 100644 --- a/pyatlan/client/common/__init__.py +++ b/pyatlan/client/common/__init__.py @@ -128,8 +128,11 @@ SSOCreateGroupMapping, SSODeleteGroupMapping, SSOGetAllGroupMappings, + SSOGetAllIdentityProviders, SSOGetGroupMapping, SSOUpdateGroupMapping, + SSOUpdateIdentityProvider, + normalize_signing_certificate, ) # Task shared logic classes @@ -284,6 +287,9 @@ # SSO shared logic classes "SSOCheckExistingMappings", "SSOCreateGroupMapping", + "SSOGetAllIdentityProviders", + "SSOUpdateIdentityProvider", + "normalize_signing_certificate", "SSODeleteGroupMapping", "SSOGetAllGroupMappings", "SSOGetGroupMapping", diff --git a/pyatlan/client/common/sso.py b/pyatlan/client/common/sso.py index 53594847f..245fa9ce8 100644 --- a/pyatlan/client/common/sso.py +++ b/pyatlan/client/common/sso.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2025 Atlan Pte. Ltd. +import base64 +import binascii from typing import List from pydantic.v1 import ValidationError, parse_obj_as @@ -8,13 +10,15 @@ from pyatlan.client.constants import ( CREATE_SSO_GROUP_MAPPING, DELETE_SSO_GROUP_MAPPING, + GET_ALL_IDPS, GET_ALL_SSO_GROUP_MAPPING, GET_SSO_GROUP_MAPPING, + UPDATE_IDP, UPDATE_SSO_GROUP_MAPPING, ) from pyatlan.errors import ErrorCode from pyatlan.model.group import AtlanGroup -from pyatlan.model.sso import SSOMapper, SSOMapperConfig +from pyatlan.model.sso import SSOMapper, SSOMapperConfig, SSOProvider from pyatlan.utils import get_epoch_timestamp GROUP_MAPPER_ATTRIBUTE = "memberOf" @@ -272,3 +276,95 @@ def check_existing_group_mappings( raise ErrorCode.SSO_GROUP_MAPPING_ALREADY_EXISTS.exception_with_parameters( atlan_group.alias, group_map.config.attribute_value ) + + +class SSOGetAllIdentityProviders: + """Shared logic for retrieving all SSO identity providers.""" + + @staticmethod + def prepare_request() -> tuple: + """ + Prepare the request for retrieving all identity providers. + + :returns: tuple of (endpoint, request_obj) + """ + return GET_ALL_IDPS, None + + @staticmethod + def process_response(raw_json) -> List[SSOProvider]: + """ + Process the raw API response into a list of identity providers. + + :param raw_json: raw API response + :returns: list of the tenant's SSO identity providers + """ + if not raw_json: + return [] + try: + return parse_obj_as(List[SSOProvider], raw_json) + except ValidationError as err: + raise ErrorCode.JSON_ERROR.exception_with_parameters( + raw_json, 200, str(err) + ) from err + + +class SSOUpdateIdentityProvider: + """Shared logic for updating an SSO identity provider.""" + + @staticmethod + def prepare_request(provider: SSOProvider) -> tuple: + """ + Prepare the request for updating an identity provider. + + The backend treats this update as a full replacement of the + provider's configuration, so `provider` must be the complete + object (retrieve it first, modify it, then pass it here) - + never a partial one, or omitted fields may be reset. + + :param provider: the complete identity provider configuration to store + :returns: tuple of (endpoint, request_obj) + """ + if not provider.alias: + raise ErrorCode.MISSING_REQUIRED_QUERY_PARAM.exception_with_parameters( + "the identity provider", "alias" + ) + endpoint = UPDATE_IDP.format_path({"sso_alias": provider.alias}) + return endpoint, provider + + +def normalize_signing_certificate(certificate: str) -> str: + """ + Convert an X.509 certificate to the form the SSO configuration stores: + a single line of base64, with no BEGIN/END lines and no whitespace. + + Accepts PEM or already-normalized input. Exactly one certificate must + be provided. + + :param certificate: certificate as PEM or single-line base64 + :raises InvalidRequestError: if the input contains more than one + certificate, no certificate, or is not base64 + :returns: single-line base64 certificate value + """ + cert_count = certificate.count("BEGIN CERTIFICATE") + if cert_count > 1: + raise ErrorCode.INVALID_CERTIFICATE.exception_with_parameters( + f"found {cert_count} certificates, expected 1 (IdP metadata files " + "often include both the old and new certificate - pass only the new one)" + ) + lines = [ + line.strip() + for line in certificate.strip().splitlines() + if line.strip() and "CERTIFICATE" not in line + ] + normalized = "".join("".join(line.split()) for line in lines) + if not normalized: + raise ErrorCode.INVALID_CERTIFICATE.exception_with_parameters( + "no certificate content between the BEGIN/END lines" + ) + try: + base64.b64decode(normalized, validate=True) + except binascii.Error as err: + raise ErrorCode.INVALID_CERTIFICATE.exception_with_parameters( + f"not base64 ({err}) - check for truncation or non-certificate text" + ) from err + return normalized diff --git a/pyatlan/client/constants.py b/pyatlan/client/constants.py index 11d7842d1..8a1b5ec2d 100644 --- a/pyatlan/client/constants.py +++ b/pyatlan/client/constants.py @@ -633,6 +633,14 @@ SSO_API = "idp/" SSO_GROUP_MAPPER = SSO_API + "{sso_alias}/mappers" +GET_ALL_IDPS = API("idp", HTTPMethod.GET, HTTPStatus.OK, endpoint=EndPoint.HERACLES) +UPDATE_IDP = API( + SSO_API + "{sso_alias}", + HTTPMethod.POST, + HTTPStatus.OK, + endpoint=EndPoint.HERACLES, +) + GET_SSO_GROUP_MAPPING = API( SSO_GROUP_MAPPER + "/{group_map_id}", HTTPMethod.GET, diff --git a/pyatlan/client/sso.py b/pyatlan/client/sso.py index 98333dbd0..73b96f2a8 100644 --- a/pyatlan/client/sso.py +++ b/pyatlan/client/sso.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, cast from pydantic.v1 import validate_arguments @@ -8,12 +8,15 @@ SSOCreateGroupMapping, SSODeleteGroupMapping, SSOGetAllGroupMappings, + SSOGetAllIdentityProviders, SSOGetGroupMapping, SSOUpdateGroupMapping, + SSOUpdateIdentityProvider, + normalize_signing_certificate, ) from pyatlan.errors import ErrorCode from pyatlan.model.group import AtlanGroup -from pyatlan.model.sso import SSOMapper +from pyatlan.model.sso import SSOMapper, SSOProvider class SSOClient: @@ -134,3 +137,81 @@ def delete_group_mapping(self, sso_alias: str, group_map_id: str) -> None: ) raw_json = self._client._call_api(endpoint, request_obj=request_obj) return raw_json + + def get_all_identity_providers(self) -> List[SSOProvider]: + """ + Retrieves all SSO identity providers configured on the tenant. + + Requires an API token with the admin role or workspace-admin + subrole (READ_TENANT_IDP); other tokens receive a 403. + + :raises AtlanError: on any error during API invocation. + :returns: list of the tenant's SSO identity providers + """ + endpoint, request_obj = SSOGetAllIdentityProviders.prepare_request() + raw_json = self._client._call_api(endpoint, request_obj=request_obj) + return SSOGetAllIdentityProviders.process_response(raw_json) + + @validate_arguments + def get_identity_provider(self, sso_alias: str) -> SSOProvider: + """ + Retrieves the SSO identity provider with the given alias. + + :param sso_alias: alias of the SSO provider (e.g. `azure`, `okta`) + :raises AtlanError: on any error during API invocation. + :raises NotFoundError: if no identity provider exists with the given alias. + :returns: the identity provider configuration + """ + for provider in self.get_all_identity_providers(): + if provider.alias == sso_alias: + return provider + raise ErrorCode.IDP_NOT_FOUND_BY_ALIAS.exception_with_parameters(sso_alias) + + @validate_arguments + def update_identity_provider(self, provider: SSOProvider) -> SSOProvider: + """ + Updates an SSO identity provider's configuration. + + The backend treats this as a full replacement: always retrieve the + current configuration first (`get_identity_provider()`), modify it, + and pass the complete object here. Sending a partial object may + silently reset fields that were omitted. + + Requires an API token with the admin role or workspace-admin + subrole (UPDATE_TENANT_IDP); other tokens receive a 403. + + :param provider: the complete identity provider configuration to store + :raises AtlanError: on any error during API invocation. + :returns: the identity provider configuration, re-read after the update + """ + endpoint, request_obj = SSOUpdateIdentityProvider.prepare_request(provider) + alias = cast(str, provider.alias) # validated non-empty by prepare_request + self._client._call_api(endpoint, request_obj=request_obj) + return self.get_identity_provider(sso_alias=alias) + + @validate_arguments + def update_signing_certificate( + self, sso_alias: str, certificate: str + ) -> SSOProvider: + """ + Replaces the signing certificate on the given SSO identity + provider, leaving the rest of the configuration untouched. + + API tokens authenticate independently of SSO, so this works even + while SSO logins are failing (for example, after the certificate + expired) - as long as the token already exists. + + :param sso_alias: alias of the SSO provider (e.g. `azure`, `okta`) + :param certificate: new X.509 certificate, as PEM or a single line + of base64; stored as one line with no BEGIN/END lines and no + line breaks + :raises AtlanError: on any error during API invocation. + :raises NotFoundError: if no identity provider exists with the given alias. + :returns: the identity provider configuration, re-read after the update + """ + provider = self.get_identity_provider(sso_alias=sso_alias) + provider.config = provider.config or {} + provider.config["signingCertificate"] = normalize_signing_certificate( + certificate + ) + return self.update_identity_provider(provider=provider) diff --git a/pyatlan/errors.py b/pyatlan/errors.py index 2a51e279d..9818eb650 100644 --- a/pyatlan/errors.py +++ b/pyatlan/errors.py @@ -284,6 +284,15 @@ class ErrorCode(Enum): "a new connection. Without at least one admin, the connection will be inaccessible to all.", InvalidRequestError, ) + INVALID_CERTIFICATE = ( + 400, + "ATLAN-PYTHON-400-080", + "Invalid signing certificate: {0}.", + "The SSO configuration stores one X.509 certificate as a single line of " + "base64, with no BEGIN/END lines and no line breaks. PEM input is " + "accepted and converted to that form.", + InvalidRequestError, + ) MISSING_PERSONA_ID = ( 400, "ATLAN-PYTHON-400-023", @@ -995,6 +1004,13 @@ class ErrorCode(Enum): "Verify the role description provided matches one of the available roles.", NotFoundError, ) + IDP_NOT_FOUND_BY_ALIAS = ( + 404, + "ATLAN-PYTHON-404-031", + "Identity provider with alias '{0}' does not exist.", + "Verify the SSO provider alias (for example: 'azure', 'okta', 'google').", + NotFoundError, + ) CONFLICT_PASSTHROUGH = ( 409, "ATLAN-PYTHON-409-000", diff --git a/pyatlan/model/sso.py b/pyatlan/model/sso.py index 215add5e7..7f570cd94 100644 --- a/pyatlan/model/sso.py +++ b/pyatlan/model/sso.py @@ -1,6 +1,6 @@ -from typing import Optional +from typing import Any, Dict, Optional -from pydantic.v1 import Field +from pydantic.v1 import Extra, Field from pyatlan.model.core import AtlanObject @@ -25,3 +25,34 @@ class SSOMapper(AtlanObject): identity_provider_mapper: str identity_provider_alias: str config: SSOMapperConfig + + +class SSOProvider(AtlanObject): + """ + A tenant's SSO identity provider configuration (Keycloak identity + provider representation), as returned by `GET /api/service/idp`. + + The nested `config` is intentionally an untyped mapping: the backend + treats updates as full replacements, so every key returned by the API + must be sent back verbatim on update. Typing it would risk silently + dropping (and therefore resetting) fields the SDK does not know about. + """ + + class Config(AtlanObject.Config): + extra = Extra.allow + + alias: Optional[str] = Field(default=None) + internal_id: Optional[str] = Field(default=None, alias="internalId") + display_name: Optional[str] = Field(default=None, alias="displayName") + provider_id: Optional[str] = Field(default=None, alias="providerId") + enabled: Optional[bool] = Field(default=None) + trust_email: Optional[bool] = Field(default=None, alias="trustEmail") + store_token: Optional[bool] = Field(default=None, alias="storeToken") + link_only: Optional[bool] = Field(default=None, alias="linkOnly") + add_read_token_role_on_create: Optional[bool] = Field( + default=None, alias="addReadTokenRoleOnCreate" + ) + first_broker_login_flow_alias: Optional[str] = Field( + default=None, alias="firstBrokerLoginFlowAlias" + ) + config: Optional[Dict[str, Any]] = Field(default=None) diff --git a/pyatlan_v9/client/aio/sso.py b/pyatlan_v9/client/aio/sso.py index c7ab80d0e..01a86ec47 100644 --- a/pyatlan_v9/client/aio/sso.py +++ b/pyatlan_v9/client/aio/sso.py @@ -8,11 +8,14 @@ import msgspec from pyatlan.client.common import AsyncApiCaller +from pyatlan.client.common.sso import normalize_signing_certificate from pyatlan.client.constants import ( CREATE_SSO_GROUP_MAPPING, DELETE_SSO_GROUP_MAPPING, + GET_ALL_IDPS, GET_ALL_SSO_GROUP_MAPPING, GET_SSO_GROUP_MAPPING, + UPDATE_IDP, UPDATE_SSO_GROUP_MAPPING, ) from pyatlan.errors import AtlanError, ErrorCode @@ -25,7 +28,7 @@ _resolve_sso_alias, ) from pyatlan_v9.model.group import AtlanGroup -from pyatlan_v9.model.sso import SSOMapper, SSOMapperConfig +from pyatlan_v9.model.sso import SSOMapper, SSOMapperConfig, SSOProvider from pyatlan_v9.validate import validate_arguments @@ -199,3 +202,95 @@ async def delete_group_mapping(self, sso_alias: str, group_map_id: str) -> None: ) raw_json = await self._client._call_api(endpoint) return raw_json + + @staticmethod + def _parse_sso_providers(raw_json) -> List[SSOProvider]: + if not raw_json: + return [] + try: + return msgspec.convert(raw_json, List[SSOProvider], strict=False) + except msgspec.ValidationError as err: + raise ErrorCode.JSON_ERROR.exception_with_parameters( + raw_json, 200, str(err) + ) from err + + async def get_all_identity_providers(self) -> List[SSOProvider]: + """ + Retrieves all SSO identity providers configured on the tenant. + + Requires an API token with the admin role or workspace-admin + subrole (READ_TENANT_IDP); other tokens receive a 403. + + :raises AtlanError: on any error during API invocation. + :returns: list of the tenant's SSO identity providers. + """ + raw_json = await self._client._call_api(GET_ALL_IDPS) + return self._parse_sso_providers(raw_json) + + @validate_arguments + async def get_identity_provider(self, sso_alias: str) -> SSOProvider: + """ + Retrieves the SSO identity provider with the given alias. + + :param sso_alias: alias of the SSO provider (e.g. `azure`, `okta`). + :raises AtlanError: on any error during API invocation. + :raises NotFoundError: if no identity provider exists with the given alias. + :returns: the identity provider configuration. + """ + for provider in await self.get_all_identity_providers(): + if provider.alias == sso_alias: + return provider + raise ErrorCode.IDP_NOT_FOUND_BY_ALIAS.exception_with_parameters(sso_alias) + + @validate_arguments + async def update_identity_provider(self, provider: SSOProvider) -> SSOProvider: + """ + Updates an SSO identity provider's configuration. + + The backend treats this as a full replacement: always retrieve the + current configuration first (`get_identity_provider()`), modify it, + and pass the complete object here. Sending a partial object may + silently reset fields that were omitted. + + Requires an API token with the admin role or workspace-admin + subrole (UPDATE_TENANT_IDP); other tokens receive a 403. + + :param provider: the complete identity provider configuration to store. + :raises AtlanError: on any error during API invocation. + :returns: the identity provider configuration, re-read after the update. + """ + if not provider.alias: + raise ErrorCode.MISSING_REQUIRED_QUERY_PARAM.exception_with_parameters( + "the identity provider", "alias" + ) + alias: str = provider.alias + endpoint = UPDATE_IDP.format_path({"sso_alias": alias}) + await self._client._call_api(endpoint, request_obj=provider) + return await self.get_identity_provider(sso_alias=alias) + + @validate_arguments + async def update_signing_certificate( + self, sso_alias: str, certificate: str + ) -> SSOProvider: + """ + Replaces the signing certificate on the given SSO identity + provider, leaving the rest of the configuration untouched. + + API tokens authenticate independently of SSO, so this works even + while SSO logins are failing (for example, after the certificate + expired) - as long as the token already exists. + + :param sso_alias: alias of the SSO provider (e.g. `azure`, `okta`). + :param certificate: new X.509 certificate, as PEM or a single line + of base64; stored as one line with no BEGIN/END lines and no + line breaks. + :raises AtlanError: on any error during API invocation. + :raises NotFoundError: if no identity provider exists with the given alias. + :returns: the identity provider configuration, re-read after the update. + """ + provider = await self.get_identity_provider(sso_alias=sso_alias) + provider.config = provider.config or {} + provider.config["signingCertificate"] = normalize_signing_certificate( + certificate + ) + return await self.update_identity_provider(provider=provider) diff --git a/pyatlan_v9/client/sso.py b/pyatlan_v9/client/sso.py index b19d2264e..34ad72671 100644 --- a/pyatlan_v9/client/sso.py +++ b/pyatlan_v9/client/sso.py @@ -9,17 +9,20 @@ import msgspec from pyatlan.client.common import ApiCaller +from pyatlan.client.common.sso import normalize_signing_certificate from pyatlan.client.constants import ( CREATE_SSO_GROUP_MAPPING, DELETE_SSO_GROUP_MAPPING, + GET_ALL_IDPS, GET_ALL_SSO_GROUP_MAPPING, GET_SSO_GROUP_MAPPING, + UPDATE_IDP, UPDATE_SSO_GROUP_MAPPING, ) from pyatlan.errors import AtlanError, ErrorCode from pyatlan.utils import get_epoch_timestamp from pyatlan_v9.model.group import AtlanGroup -from pyatlan_v9.model.sso import SSOMapper, SSOMapperConfig +from pyatlan_v9.model.sso import SSOMapper, SSOMapperConfig, SSOProvider from pyatlan_v9.validate import validate_arguments GROUP_MAPPER_ATTRIBUTE = "memberOf" @@ -223,3 +226,95 @@ def delete_group_mapping(self, sso_alias: str, group_map_id: str) -> None: ) raw_json = self._client._call_api(endpoint) return raw_json + + @staticmethod + def _parse_sso_providers(raw_json) -> List[SSOProvider]: + if not raw_json: + return [] + try: + return msgspec.convert(raw_json, List[SSOProvider], strict=False) + except msgspec.ValidationError as err: + raise ErrorCode.JSON_ERROR.exception_with_parameters( + raw_json, 200, str(err) + ) from err + + def get_all_identity_providers(self) -> List[SSOProvider]: + """ + Retrieves all SSO identity providers configured on the tenant. + + Requires an API token with the admin role or workspace-admin + subrole (READ_TENANT_IDP); other tokens receive a 403. + + :raises AtlanError: on any error during API invocation. + :returns: list of the tenant's SSO identity providers. + """ + raw_json = self._client._call_api(GET_ALL_IDPS) + return self._parse_sso_providers(raw_json) + + @validate_arguments + def get_identity_provider(self, sso_alias: str) -> SSOProvider: + """ + Retrieves the SSO identity provider with the given alias. + + :param sso_alias: alias of the SSO provider (e.g. `azure`, `okta`). + :raises AtlanError: on any error during API invocation. + :raises NotFoundError: if no identity provider exists with the given alias. + :returns: the identity provider configuration. + """ + for provider in self.get_all_identity_providers(): + if provider.alias == sso_alias: + return provider + raise ErrorCode.IDP_NOT_FOUND_BY_ALIAS.exception_with_parameters(sso_alias) + + @validate_arguments + def update_identity_provider(self, provider: SSOProvider) -> SSOProvider: + """ + Updates an SSO identity provider's configuration. + + The backend treats this as a full replacement: always retrieve the + current configuration first (`get_identity_provider()`), modify it, + and pass the complete object here. Sending a partial object may + silently reset fields that were omitted. + + Requires an API token with the admin role or workspace-admin + subrole (UPDATE_TENANT_IDP); other tokens receive a 403. + + :param provider: the complete identity provider configuration to store. + :raises AtlanError: on any error during API invocation. + :returns: the identity provider configuration, re-read after the update. + """ + if not provider.alias: + raise ErrorCode.MISSING_REQUIRED_QUERY_PARAM.exception_with_parameters( + "the identity provider", "alias" + ) + alias: str = provider.alias + endpoint = UPDATE_IDP.format_path({"sso_alias": alias}) + self._client._call_api(endpoint, request_obj=provider) + return self.get_identity_provider(sso_alias=alias) + + @validate_arguments + def update_signing_certificate( + self, sso_alias: str, certificate: str + ) -> SSOProvider: + """ + Replaces the signing certificate on the given SSO identity + provider, leaving the rest of the configuration untouched. + + API tokens authenticate independently of SSO, so this works even + while SSO logins are failing (for example, after the certificate + expired) - as long as the token already exists. + + :param sso_alias: alias of the SSO provider (e.g. `azure`, `okta`). + :param certificate: new X.509 certificate, as PEM or a single line + of base64; stored as one line with no BEGIN/END lines and no + line breaks. + :raises AtlanError: on any error during API invocation. + :raises NotFoundError: if no identity provider exists with the given alias. + :returns: the identity provider configuration, re-read after the update. + """ + provider = self.get_identity_provider(sso_alias=sso_alias) + provider.config = provider.config or {} + provider.config["signingCertificate"] = normalize_signing_certificate( + certificate + ) + return self.update_identity_provider(provider=provider) diff --git a/pyatlan_v9/model/sso.py b/pyatlan_v9/model/sso.py index 1724d2e90..5272a42fa 100644 --- a/pyatlan_v9/model/sso.py +++ b/pyatlan_v9/model/sso.py @@ -42,3 +42,39 @@ class SSOMapper(msgspec.Struct, kw_only=True, rename="camel", omit_defaults=True def to_dict(self) -> dict: """Serialize to dict, excluding fields with None/default values.""" return json.loads(msgspec.json.encode(self)) + + +class SSOProvider(msgspec.Struct, kw_only=True, rename="camel", omit_defaults=True): + """ + A tenant's SSO identity provider configuration (Keycloak identity + provider representation), as returned by `GET /api/service/idp`. + + The nested `config` is an untyped mapping so that every key returned + by the API is sent back verbatim on update - the backend treats + updates as full replacements, and omitted keys may be reset. + + Note: unlike the pydantic model, msgspec cannot capture unknown + top-level fields; if the API adds new top-level fields they must be + added here to survive a get-then-update round-trip. + """ + + alias: Union[str, None] = None + internal_id: Union[str, None] = None + display_name: Union[str, None] = None + provider_id: Union[str, None] = None + enabled: Union[bool, None] = None + trust_email: Union[bool, None] = None + store_token: Union[bool, None] = None + link_only: Union[bool, None] = None + add_read_token_role_on_create: Union[bool, None] = None + first_broker_login_flow_alias: Union[str, None] = None + post_broker_login_flow_alias: Union[str, None] = None + authenticate_by_default: Union[bool, None] = None + update_profile_first_login_mode: Union[str, None] = None + hide_on_login: Union[bool, None] = None + organization_id: Union[str, None] = None + config: Union[dict, None] = None + + def to_dict(self) -> dict: + """Serialize to dict, excluding fields with None/default values.""" + return json.loads(msgspec.json.encode(self)) diff --git a/tests/unit/aio/test_sso_identity_provider.py b/tests/unit/aio/test_sso_identity_provider.py new file mode 100644 index 000000000..392c1b456 --- /dev/null +++ b/tests/unit/aio/test_sso_identity_provider.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 Atlan Pte. Ltd. +from json import load +from pathlib import Path +from unittest.mock import Mock + +import pytest + +from pyatlan.client.aio.sso import AsyncSSOClient +from pyatlan.client.common import AsyncApiCaller +from pyatlan.errors import NotFoundError + +TEST_DATA_DIR = Path(__file__).parent.parent / "data" +SSO_RESPONSES_DIR = TEST_DATA_DIR / "sso_responses" + +PEM_CERT = ( + "-----BEGIN CERTIFICATE-----\n" + "TkVXQ0VS\nVA==\n" + "-----END CERTIFICATE-----\n" +) # body is base64("NEWCERT") + + +@pytest.fixture(autouse=True) +def set_env(monkeypatch): + monkeypatch.setenv("ATLAN_BASE_URL", "https://test.atlan.com") + monkeypatch.setenv("ATLAN_API_KEY", "test-api-key") + + +@pytest.fixture() +def mock_api_caller(): + return Mock(spec=AsyncApiCaller) + + +@pytest.fixture() +def get_all_idps_json(): + with (SSO_RESPONSES_DIR / "get_all_identity_providers.json").open() as f: + return load(f) + + +async def test_get_all(mock_api_caller, get_all_idps_json): + mock_api_caller._call_api.return_value = get_all_idps_json + client = AsyncSSOClient(mock_api_caller) + providers = await client.get_all_identity_providers() + assert len(providers) == 1 and providers[0].alias == "okta" + + +async def test_get_by_alias_not_found(mock_api_caller, get_all_idps_json): + mock_api_caller._call_api.return_value = get_all_idps_json + client = AsyncSSOClient(mock_api_caller) + with pytest.raises(NotFoundError): + await client.get_identity_provider("azure") + + +async def test_update_signing_certificate_round_trip( + mock_api_caller, get_all_idps_json +): + mock_api_caller._call_api.side_effect = [ + get_all_idps_json, + None, + get_all_idps_json, + ] + client = AsyncSSOClient(mock_api_caller) + await client.update_signing_certificate(sso_alias="okta", certificate=PEM_CERT) + update_call = mock_api_caller._call_api.call_args_list[1] + sent = update_call.kwargs.get("request_obj") or update_call.args[1] + assert sent.config["signingCertificate"] == "TkVXQ0VSVA==" + assert sent.config["someFutureConfigKey"] == "also-must-survive" diff --git a/tests/unit/data/sso_responses/get_all_identity_providers.json b/tests/unit/data/sso_responses/get_all_identity_providers.json new file mode 100644 index 000000000..2c69418e3 --- /dev/null +++ b/tests/unit/data/sso_responses/get_all_identity_providers.json @@ -0,0 +1,26 @@ +[ + { + "alias": "okta", + "internalId": "8f0f7a9d-0000-0000-0000-000000000000", + "displayName": "Sign in with Okta", + "providerId": "saml", + "enabled": true, + "trustEmail": true, + "storeToken": false, + "linkOnly": false, + "addReadTokenRoleOnCreate": false, + "firstBrokerLoginFlowAlias": "SimpleLoginFlow", + "futureUnknownField": "must-survive-round-trip", + "config": { + "nameIDPolicyFormat": "urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress", + "postBindingAuthnRequest": "true", + "postBindingResponse": "true", + "principalType": "SUBJECT", + "signingCertificate": "MIIDOLDCERTAAAA", + "singleSignOnServiceUrl": "https://example.okta.com/app/x/sso/saml", + "syncMode": "IMPORT", + "validateSignature": "true", + "someFutureConfigKey": "also-must-survive" + } + } +] diff --git a/tests/unit/test_sso_identity_provider.py b/tests/unit/test_sso_identity_provider.py new file mode 100644 index 000000000..e186542b6 --- /dev/null +++ b/tests/unit/test_sso_identity_provider.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 Atlan Pte. Ltd. +from json import load, loads +from pathlib import Path +from unittest.mock import Mock + +import pytest + +from pyatlan.client.common import ApiCaller, normalize_signing_certificate +from pyatlan.client.sso import SSOClient +from pyatlan.errors import InvalidRequestError, NotFoundError +from pyatlan.model.sso import SSOProvider + +TEST_DATA_DIR = Path(__file__).parent / "data" +SSO_RESPONSES_DIR = TEST_DATA_DIR / "sso_responses" +GET_ALL_IDPS_JSON = "get_all_identity_providers.json" + +PEM_CERT = ( + "-----BEGIN CERTIFICATE-----\n" + "TkVXQ0VS\nVA==\n" + "-----END CERTIFICATE-----\n" +) # body is base64("NEWCERT") + + +def load_json(filename): + with (SSO_RESPONSES_DIR / filename).open() as input_file: + return load(input_file) + + +@pytest.fixture(autouse=True) +def set_env(monkeypatch): + monkeypatch.setenv("ATLAN_BASE_URL", "https://test.atlan.com") + monkeypatch.setenv("ATLAN_API_KEY", "test-api-key") + + +@pytest.fixture() +def mock_api_caller(): + return Mock(spec=ApiCaller) + + +@pytest.fixture() +def get_all_idps_json(): + return load_json(GET_ALL_IDPS_JSON) + + +class TestNormalizeSigningCertificate: + def test_pem_input_becomes_single_line(self): + assert normalize_signing_certificate(PEM_CERT) == "TkVXQ0VSVA==" + + def test_raw_single_line_unchanged(self): + assert normalize_signing_certificate("TkVXQ0VSVA==") == "TkVXQ0VSVA==" + + def test_crlf_and_blank_lines_stripped(self): + cert = "\r\nTkVXQ0VS\r\nVA==\r\n\r\n" + assert normalize_signing_certificate(cert) == "TkVXQ0VSVA==" + + def test_internal_whitespace_stripped(self): + assert normalize_signing_certificate("TkVXQ0VS \tVA==") == "TkVXQ0VSVA==" + + def test_multi_cert_bundle_rejected(self): + """IdP federation metadata often bundles the expiring cert AND its + replacement; silently concatenating them would write an invalid + certificate into the tenant's SSO config during an outage.""" + two = PEM_CERT + PEM_CERT + with pytest.raises( + InvalidRequestError, match="found 2 certificates, expected 1" + ): + normalize_signing_certificate(two) + + def test_non_base64_rejected(self): + with pytest.raises(InvalidRequestError, match="not base64"): + normalize_signing_certificate("not a certificate at all!!!") + + def test_empty_input_rejected(self): + with pytest.raises(InvalidRequestError, match="no certificate content between"): + normalize_signing_certificate( + "-----BEGIN CERTIFICATE-----\n-----END CERTIFICATE-----" + ) + + +class TestGetIdentityProviders: + def test_get_all(self, mock_api_caller, get_all_idps_json): + mock_api_caller._call_api.return_value = get_all_idps_json + client = SSOClient(mock_api_caller) + providers = client.get_all_identity_providers() + assert len(providers) == 1 + assert providers[0].alias == "okta" + assert providers[0].provider_id == "saml" + assert providers[0].config["signingCertificate"] == "MIIDOLDCERTAAAA" + + def test_get_all_empty(self, mock_api_caller): + mock_api_caller._call_api.return_value = [] + client = SSOClient(mock_api_caller) + assert client.get_all_identity_providers() == [] + + def test_get_by_alias(self, mock_api_caller, get_all_idps_json): + mock_api_caller._call_api.return_value = get_all_idps_json + client = SSOClient(mock_api_caller) + provider = client.get_identity_provider("okta") + assert provider.alias == "okta" + + def test_get_by_alias_not_found(self, mock_api_caller, get_all_idps_json): + mock_api_caller._call_api.return_value = get_all_idps_json + client = SSOClient(mock_api_caller) + with pytest.raises(NotFoundError): + client.get_identity_provider("azure") + + +class TestUpdateIdentityProvider: + def test_full_object_round_trip_preserves_all_fields( + self, mock_api_caller, get_all_idps_json + ): + """Regression guard for the partial-update trap (SHA-497 / BLDX-634): + every field returned by the API — including ones this SDK version + does not know about — must be present in the update payload.""" + mock_api_caller._call_api.return_value = get_all_idps_json + client = SSOClient(mock_api_caller) + provider = client.get_identity_provider("okta") + + payload = loads(provider.json(by_alias=True, exclude_none=True)) + source = get_all_idps_json[0] + for key in source["config"]: + assert key in payload["config"], f"config key dropped: {key}" + assert payload["config"]["someFutureConfigKey"] == "also-must-survive" + # top-level unknown fields must also survive (Extra.allow), and the + # internal extras holder must never leak into the payload + assert payload["futureUnknownField"] == "must-survive-round-trip" + assert "__atlan_extra__" not in payload + + def test_update_requires_alias(self, mock_api_caller): + client = SSOClient(mock_api_caller) + provider = SSOProvider(config={"signingCertificate": "MIIDOLDCERTAAAA"}) + with pytest.raises(InvalidRequestError): + client.update_identity_provider(provider=provider) + + def test_update_posts_to_alias_path_and_refetches( + self, mock_api_caller, get_all_idps_json + ): + # first call: GET (for get_identity_provider); second: POST update; + # third: GET (re-read after update) + mock_api_caller._call_api.side_effect = [ + get_all_idps_json, + None, + get_all_idps_json, + ] + client = SSOClient(mock_api_caller) + provider = client.get_identity_provider("okta") + result = client.update_identity_provider(provider=provider) + assert result.alias == "okta" + update_call = mock_api_caller._call_api.call_args_list[1] + endpoint = ( + update_call.args[0] if update_call.args else update_call.kwargs["api"] + ) + assert "idp/okta" in endpoint.path + + def test_update_signing_certificate_normalizes_and_preserves_config( + self, mock_api_caller, get_all_idps_json + ): + mock_api_caller._call_api.side_effect = [ + get_all_idps_json, + None, + get_all_idps_json, + ] + client = SSOClient(mock_api_caller) + client.update_signing_certificate(sso_alias="okta", certificate=PEM_CERT) + update_call = mock_api_caller._call_api.call_args_list[1] + sent = update_call.kwargs.get("request_obj") or update_call.args[1] + assert sent.config["signingCertificate"] == "TkVXQ0VSVA==" + # every other config key untouched + assert sent.config["singleSignOnServiceUrl"] == ( + "https://example.okta.com/app/x/sso/saml" + ) + assert sent.config["someFutureConfigKey"] == "also-must-survive" diff --git a/tests_v9/unit/aio/test_sso_identity_provider.py b/tests_v9/unit/aio/test_sso_identity_provider.py new file mode 100644 index 000000000..634e0c51b --- /dev/null +++ b/tests_v9/unit/aio/test_sso_identity_provider.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 Atlan Pte. Ltd. +from json import load +from pathlib import Path +from unittest.mock import Mock + +import pytest + +from pyatlan.client.common import AsyncApiCaller +from pyatlan.errors import NotFoundError +from pyatlan_v9.client.aio.sso import V9AsyncSSOClient as AsyncSSOClient + +TEST_DATA_DIR = Path(__file__).parent.parent.parent.parent / "tests" / "unit" / "data" +SSO_RESPONSES_DIR = TEST_DATA_DIR / "sso_responses" + +PEM_CERT = ( + "-----BEGIN CERTIFICATE-----\n" + "TkVXQ0VS\nVA==\n" + "-----END CERTIFICATE-----\n" +) # body is base64("NEWCERT") + + +@pytest.fixture(autouse=True) +def set_env(monkeypatch): + monkeypatch.setenv("ATLAN_BASE_URL", "https://test.atlan.com") + monkeypatch.setenv("ATLAN_API_KEY", "test-api-key") + + +@pytest.fixture() +def mock_api_caller(): + return Mock(spec=AsyncApiCaller) + + +@pytest.fixture() +def get_all_idps_json(): + with (SSO_RESPONSES_DIR / "get_all_identity_providers.json").open() as f: + return load(f) + + +async def test_get_all(mock_api_caller, get_all_idps_json): + mock_api_caller._call_api.return_value = get_all_idps_json + client = AsyncSSOClient(mock_api_caller) + providers = await client.get_all_identity_providers() + assert len(providers) == 1 and providers[0].alias == "okta" + + +async def test_get_by_alias_not_found(mock_api_caller, get_all_idps_json): + mock_api_caller._call_api.return_value = get_all_idps_json + client = AsyncSSOClient(mock_api_caller) + with pytest.raises(NotFoundError): + await client.get_identity_provider("azure") + + +async def test_update_signing_certificate_round_trip( + mock_api_caller, get_all_idps_json +): + mock_api_caller._call_api.side_effect = [ + get_all_idps_json, + None, + get_all_idps_json, + ] + client = AsyncSSOClient(mock_api_caller) + await client.update_signing_certificate(sso_alias="okta", certificate=PEM_CERT) + update_call = mock_api_caller._call_api.call_args_list[1] + sent = update_call.kwargs.get("request_obj") or update_call.args[1] + assert sent.config["signingCertificate"] == "TkVXQ0VSVA==" + assert sent.config["someFutureConfigKey"] == "also-must-survive" diff --git a/tests_v9/unit/test_sso_identity_provider.py b/tests_v9/unit/test_sso_identity_provider.py new file mode 100644 index 000000000..e0a2f41c7 --- /dev/null +++ b/tests_v9/unit/test_sso_identity_provider.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 Atlan Pte. Ltd. + +""" +Unit tests for v9 SSO identity provider methods — ported from +tests/unit/test_sso_identity_provider.py. +""" + +from json import load +from pathlib import Path +from unittest.mock import Mock + +import pytest + +from pyatlan.client.common import ApiCaller +from pyatlan.errors import InvalidRequestError, NotFoundError +from pyatlan_v9.client.sso import V9SSOClient as SSOClient +from pyatlan_v9.model.sso import SSOProvider + +TEST_DATA_DIR = Path(__file__).parent.parent.parent / "tests" / "unit" / "data" +SSO_RESPONSES_DIR = TEST_DATA_DIR / "sso_responses" +GET_ALL_IDPS_JSON = "get_all_identity_providers.json" + +PEM_CERT = ( + "-----BEGIN CERTIFICATE-----\n" + "TkVXQ0VS\nVA==\n" + "-----END CERTIFICATE-----\n" +) # body is base64("NEWCERT") + + +def load_json(filename): + with (SSO_RESPONSES_DIR / filename).open() as input_file: + return load(input_file) + + +@pytest.fixture(autouse=True) +def set_env(monkeypatch): + monkeypatch.setenv("ATLAN_BASE_URL", "https://test.atlan.com") + monkeypatch.setenv("ATLAN_API_KEY", "test-api-key") + + +@pytest.fixture() +def mock_api_caller(): + return Mock(spec=ApiCaller) + + +@pytest.fixture() +def get_all_idps_json(): + return load_json(GET_ALL_IDPS_JSON) + + +class TestGetIdentityProviders: + def test_get_all(self, mock_api_caller, get_all_idps_json): + mock_api_caller._call_api.return_value = get_all_idps_json + client = SSOClient(mock_api_caller) + providers = client.get_all_identity_providers() + assert len(providers) == 1 + assert isinstance(providers[0], SSOProvider) + assert providers[0].alias == "okta" + assert providers[0].provider_id == "saml" + assert providers[0].config["signingCertificate"] == "MIIDOLDCERTAAAA" + + def test_get_all_empty(self, mock_api_caller): + mock_api_caller._call_api.return_value = [] + client = SSOClient(mock_api_caller) + assert client.get_all_identity_providers() == [] + + def test_get_by_alias_not_found(self, mock_api_caller, get_all_idps_json): + mock_api_caller._call_api.return_value = get_all_idps_json + client = SSOClient(mock_api_caller) + with pytest.raises(NotFoundError): + client.get_identity_provider("azure") + + +class TestUpdateIdentityProvider: + def test_config_keys_survive_round_trip(self, mock_api_caller, get_all_idps_json): + """Regression guard (SHA-497 / BLDX-634): every `config` key returned + by the API — including keys this SDK version does not know about — + must be present in the serialized update payload. + + Note: msgspec drops unknown TOP-LEVEL fields (unlike the pydantic + model, which uses Extra.allow); this is a known limitation pinned + in the model docstring.""" + mock_api_caller._call_api.return_value = get_all_idps_json + client = SSOClient(mock_api_caller) + provider = client.get_identity_provider("okta") + payload = provider.to_dict() + source = get_all_idps_json[0] + for key in source["config"]: + assert key in payload["config"], f"config key dropped: {key}" + assert payload["config"]["someFutureConfigKey"] == "also-must-survive" + + def test_update_requires_alias(self, mock_api_caller): + client = SSOClient(mock_api_caller) + provider = SSOProvider(config={"signingCertificate": "MIIDOLDCERTAAAA"}) + with pytest.raises(InvalidRequestError): + client.update_identity_provider(provider=provider) + + def test_update_posts_to_alias_path_and_refetches( + self, mock_api_caller, get_all_idps_json + ): + mock_api_caller._call_api.side_effect = [ + get_all_idps_json, + None, + get_all_idps_json, + ] + client = SSOClient(mock_api_caller) + provider = client.get_identity_provider("okta") + result = client.update_identity_provider(provider=provider) + assert result.alias == "okta" + update_call = mock_api_caller._call_api.call_args_list[1] + endpoint = update_call.args[0] + assert "idp/okta" in endpoint.path + + def test_update_signing_certificate_normalizes_and_preserves_config( + self, mock_api_caller, get_all_idps_json + ): + mock_api_caller._call_api.side_effect = [ + get_all_idps_json, + None, + get_all_idps_json, + ] + client = SSOClient(mock_api_caller) + client.update_signing_certificate(sso_alias="okta", certificate=PEM_CERT) + update_call = mock_api_caller._call_api.call_args_list[1] + sent = update_call.kwargs.get("request_obj") or update_call.args[1] + assert sent.config["signingCertificate"] == "TkVXQ0VSVA==" + assert sent.config["someFutureConfigKey"] == "also-must-survive"