fixing flags (WIP)

Author: jp-agenta

api/ee/src/apis/fastapi/organizations/router.py

@@ -16,6 +16,7 @@
     DomainVerificationService,
     SSOProviderService,
 )
+from ee.src.services import db_manager_ee
 from ee.src.utils.permissions import check_user_org_access
 from ee.src.services.selectors import get_user_org_and_workspace_id
 
@@ -35,6 +36,36 @@ async def verify_user_org_access(user_id: str, organization_id: str) -> None:
         )
 
 
+async def require_email_or_social_or_root_enabled(organization_id: str) -> None:
+    """Block domain/provider changes when SSO is the only allowed method."""
+    organization = await db_manager_ee.get_organization(organization_id)
+    flags = organization.flags or {}
+    allow_email = flags.get("allow_email", False)
+    allow_social = flags.get("allow_social", False)
+    allow_root = flags.get("allow_root", False)
+    if not (allow_email or allow_social or allow_root):
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "To modify domains or SSO providers, enable email or social authentication "
+                "for this organization, or enable root access for owners."
+            ),
+        )
+
+
+async def require_domains_and_auto_join_disabled(organization_id: str) -> None:
+    """Block edits to verified domains when domains-only or auto-join is enabled."""
+    organization = await db_manager_ee.get_organization(organization_id)
+    flags = organization.flags or {}
+    if flags.get("domains_only") or flags.get("auto_join"):
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "Disable domains-only and auto-join before modifying verified domains."
+            ),
+        )
+
+
 # Domain Verification Endpoints
 
 
@@ -81,6 +112,7 @@ async def verify_domain(
     user_id = request.state.user_id
 
     await verify_user_org_access(user_id, organization_id)
+    await require_domains_and_auto_join_disabled(organization_id)
 
     return await domain_service.verify_domain(
         organization_id, payload.domain_id, user_id
@@ -115,6 +147,7 @@ async def refresh_domain_token(
     user_id = request.state.user_id
 
     await verify_user_org_access(user_id, organization_id)
+    await require_domains_and_auto_join_disabled(organization_id)
 
     return await domain_service.refresh_token(organization_id, domain_id, user_id)
 
@@ -134,6 +167,7 @@ async def reset_domain(
     user_id = request.state.user_id
 
     await verify_user_org_access(user_id, organization_id)
+    await require_domains_and_auto_join_disabled(organization_id)
 
     return await domain_service.reset_domain(organization_id, domain_id, user_id)
 
@@ -148,6 +182,7 @@ async def delete_domain(
     user_id = request.state.user_id
 
     await verify_user_org_access(user_id, organization_id)
+    await require_domains_and_auto_join_disabled(organization_id)
 
     await domain_service.delete_domain(organization_id, domain_id, user_id)
     return Response(status_code=204)
@@ -172,6 +207,7 @@ async def create_provider(
     user_id = request.state.user_id
 
     await verify_user_org_access(user_id, organization_id)
+    await require_email_or_social_or_root_enabled(organization_id)
 
     return await provider_service.create_provider(organization_id, payload, user_id)
 
@@ -187,6 +223,7 @@ async def update_provider(
     user_id = request.state.user_id
 
     await verify_user_org_access(user_id, organization_id)
+    await require_email_or_social_or_root_enabled(organization_id)
 
     return await provider_service.update_provider(
         organization_id, provider_id, payload, user_id
@@ -225,6 +262,7 @@ async def test_provider(
     user_id = request.state.user_id
 
     await verify_user_org_access(user_id, organization_id)
+    await require_email_or_social_or_root_enabled(organization_id)
 
     return await provider_service.test_provider(organization_id, provider_id, user_id)
 
@@ -239,6 +277,7 @@ async def delete_provider(
     user_id = request.state.user_id
 
     await verify_user_org_access(user_id, organization_id)
+    await require_email_or_social_or_root_enabled(organization_id)
 
     await provider_service.delete_provider(organization_id, provider_id, user_id)
     return Response(status_code=204)

api/ee/src/services/db_manager_ee.py

@@ -42,6 +42,10 @@
     UserDB,
     InvitationDB,
 )
+from ee.src.dbs.postgres.organizations.dao import (
+    OrganizationProvidersDAO,
+    OrganizationDomainsDAO,
+)
 from ee.src.services.converters import get_workspace_in_format
 from ee.src.services.selectors import get_org_default_workspace
 
@@ -1133,6 +1137,55 @@ async def update_organization(
                 allow_sso = merged_flags.get("allow_sso", False)
                 allow_root = merged_flags.get("allow_root", False)
 
+                changing_auth_flags = any(
+                    key in new_flags for key in ("allow_email", "allow_social", "allow_sso")
+                )
+                changing_auto_join = "auto_join" in new_flags
+                changing_domains_only = "domains_only" in new_flags
+
+                if changing_auth_flags and allow_sso:
+                    providers_dao = OrganizationProvidersDAO(session)
+                    providers = await providers_dao.list_by_organization(organization_id)
+                    active_valid = [
+                        provider
+                        for provider in providers
+                        if (provider.flags or {}).get("is_active")
+                        and (provider.flags or {}).get("is_valid")
+                    ]
+                    if not active_valid:
+                        raise ValueError(
+                            "SSO cannot be enabled until at least one SSO provider is "
+                            "active and verified."
+                        )
+                    if not allow_email and not allow_social:
+                        if not active_valid:
+                            raise ValueError(
+                                "SSO-only authentication requires at least one SSO provider to "
+                                "be active and verified."
+                            )
+
+                if changing_auto_join and merged_flags.get("auto_join", False):
+                    domains_dao = OrganizationDomainsDAO(session)
+                    domains = await domains_dao.list_by_organization(organization_id)
+                    has_verified_domain = any(
+                        (domain.flags or {}).get("is_verified") for domain in domains
+                    )
+                    if not has_verified_domain:
+                        raise ValueError(
+                            "Auto-join requires at least one verified domain."
+                        )
+
+                if changing_domains_only and merged_flags.get("domains_only", False):
+                    domains_dao = OrganizationDomainsDAO(session)
+                    domains = await domains_dao.list_by_organization(organization_id)
+                    has_verified_domain = any(
+                        (domain.flags or {}).get("is_verified") for domain in domains
+                    )
+                    if not has_verified_domain:
+                        raise ValueError(
+                            "Domains-only requires at least one verified domain."
+                        )
+
                 # Check if all auth methods are disabled
                 all_auth_disabled = not (allow_email or allow_social or allow_sso)
 

api/ee/src/services/organization_security_service.py

@@ -30,6 +30,7 @@
     OrganizationProviderUpdate,
     OrganizationProviderResponse,
 )
+from ee.src.services import db_manager_ee
 
 logger = logging.getLogger(__name__)
 
@@ -660,6 +661,26 @@ async def delete_provider(
             if not provider:
                 raise HTTPException(status_code=404, detail="Provider not found")
 
+            organization = await db_manager_ee.get_organization(organization_id)
+            flags = organization.flags or {}
+            if flags.get("allow_sso"):
+                providers = await dao.list_by_organization(organization_id)
+                remaining = [
+                    p
+                    for p in providers
+                    if str(p.id) != str(provider_id)
+                    and (p.flags or {}).get("is_active")
+                    and (p.flags or {}).get("is_valid")
+                ]
+                if not remaining:
+                    raise HTTPException(
+                        status_code=400,
+                        detail=(
+                            "Cannot delete the last active and verified SSO provider while "
+                            "SSO is enabled."
+                        ),
+                    )
+
             await self._vault_service().delete_secret(
                 secret_id=provider.secret_id,
                 organization_id=organization_id,

api/ee/src/services/workspace_manager.py

@@ -30,6 +30,7 @@
     check_valid_invitation,
 )
 from ee.src.services.organization_service import send_invitation_email
+from ee.src.dbs.postgres.organizations.dao import OrganizationDomainsDAO
 
 log = get_module_logger(__name__)
 
@@ -155,6 +156,30 @@ async def invite_user_to_workspace(
         organization = await db_manager_ee.get_organization(organization_id)
         user_performing_action = await db_manager.get_user(user_uid)
 
+        # Check if domains_only is enabled for this organization
+        org_flags = organization.flags or {}
+        domains_only = org_flags.get("domains_only", False)
+
+        # If domains_only is enabled, get the list of verified domains
+        verified_domain_slugs = set()
+        if domains_only:
+            domains_dao = OrganizationDomainsDAO()
+            org_domains = await domains_dao.list_by_organization(organization_id)
+            verified_domain_slugs = {
+                d.slug.lower()
+                for d in org_domains
+                if d.flags and d.flags.get("is_verified", False)
+            }
+
+            # If domains_only is enabled but no verified domains exist, block all invitations
+            if not verified_domain_slugs:
+                return JSONResponse(
+                    status_code=400,
+                    content={
+                        "error": "Cannot send invitations: domains_only is enabled but no verified domains exist"
+                    },
+                )
+
         for payload_invite in payload:
             # Check that the user is not inviting themselves
             if payload_invite.email == user_performing_action.email:
@@ -163,6 +188,17 @@ async def invite_user_to_workspace(
                     content={"error": "You cannot invite yourself to a workspace"},
                 )
 
+            # Check if domains_only is enabled and validate the email domain
+            if domains_only:
+                email_domain = payload_invite.email.split("@")[-1].lower()
+                if email_domain not in verified_domain_slugs:
+                    return JSONResponse(
+                        status_code=400,
+                        content={
+                            "error": f"Cannot invite {payload_invite.email}: domain '{email_domain}' is not a verified domain for this organization"
+                        },
+                    )
+
             # Check if the user is already a member of the workspace
             if await db_manager_ee.check_user_in_workspace_with_email(
                 payload_invite.email, str(workspace.id)

api/oss/src/apis/fastapi/auth/router.py

@@ -74,7 +74,7 @@ async def check_organization_access(request: Request, organization_id: str):
 
     if policy_error and policy_error.get("error") in {
         "AUTH_UPGRADE_REQUIRED",
-        "AUTH_SSO_DISABLED",
+        "AUTH_SSO_DENIED",
     }:
         detail = {
             "error": policy_error.get("error"),
@@ -90,9 +90,7 @@ async def check_organization_access(request: Request, organization_id: str):
 
 
 @auth_router.post("/session/identities")
-async def update_session_identities(
-    request: Request, payload: SessionIdentitiesUpdate
-):
+async def update_session_identities(request: Request, payload: SessionIdentitiesUpdate):
     try:
         session = await get_session(request)  # type: ignore
     except Exception:
@@ -117,7 +115,9 @@ async def update_session_identities(
     elif hasattr(session, "merge_into_access_token_payload"):
         await session.merge_into_access_token_payload({"session_identities": merged})
     else:
-        raise HTTPException(status_code=500, detail="Session payload update not supported")
+        raise HTTPException(
+            status_code=500, detail="Session payload update not supported"
+        )
     return {"session_identities": merged, "previous": current}
 
 

api/oss/src/core/auth/service.py

@@ -532,15 +532,9 @@ async def enforce_domain_policies(self, email: str, user_id: UUID) -> None:
                 print(f"Error during auto-join: {e}")
 
         # 2. Domains-only enforcement: Check if user has access
-        # This is enforced at the organization level, not during login
-        # It's checked when user tries to access organization resources
-        # For now, we just validate that the domain matches
-        domains_only = org_flags.get("domains_only", False)
-        if domains_only:
-            # If domains_only is enabled, user MUST have matching domain
-            # Since we already verified the domain matches (domain_dto exists),
-            # the user is allowed. If domain didn't match, domain_dto would be None.
-            pass
+        # This is enforced at the organization level via check_organization_access()
+        # when the user tries to access organization resources through the middleware.
+        # No action needed here during login - enforcement happens at access time.
 
     # ============================================================================
     # AUTHORIZATION: Validate access based on policies
@@ -619,19 +613,27 @@ async def check_organization_access(
             # If the session used SSO but the org doesn't allow it (or provider inactive),
             # block and instruct user to re-auth with allowed methods.
             sso_identity = next(
-                (identity for identity in session_identities if identity.startswith("sso:")),
+                (
+                    identity
+                    for identity in session_identities
+                    if identity.startswith("sso:")
+                ),
                 None,
             )
             if sso_identity and self.providers_dao:
                 org_slug = await self._get_organization_slug(organization_id)
                 provider_slug = (
-                    sso_identity.split(":")[2] if len(sso_identity.split(":")) > 2 else None
+                    sso_identity.split(":")[2]
+                    if len(sso_identity.split(":")) > 2
+                    else None
                 )
                 providers = await self.providers_dao.list_by_organization(
                     str(organization_id)
                 )
                 active_provider_slugs = {
-                    p.slug for p in providers if p.flags and p.flags.get("is_active", False)
+                    p.slug
+                    for p in providers
+                    if p.flags and p.flags.get("is_active", False)
                 }
                 sso_matches_org = bool(
                     org_slug and sso_identity.startswith(f"sso:{org_slug}:")
@@ -647,8 +649,8 @@ async def check_organization_access(
                     if allow_social:
                         required_methods.append("social:*")
                     return {
-                        "error": "AUTH_SSO_DISABLED",
-                        "message": "SSO is not enabled for this organization",
+                        "error": "AUTH_SSO_DENIED",
+                        "message": "SSO is denied for this organization",
                         "required_methods": required_methods,
                         "current_identities": session_identities,
                     }
@@ -663,6 +665,33 @@ async def check_organization_access(
                 "sso_providers": sso_providers,
             }
 
+        # Check domains_only enforcement
+        domains_only = org_flags.get("domains_only", False)
+        if domains_only and self.domains_dao:
+            # Get user's email to check domain
+            user = await db_manager.get_user(str(user_id))
+            if user and user.email:
+                email_domain = user.email.split("@")[-1].lower()
+
+                # Get verified domains for this organization
+                org_domains = await self.domains_dao.list_by_organization(
+                    str(organization_id)
+                )
+                verified_domain_slugs = {
+                    d.slug.lower()
+                    for d in org_domains
+                    if d.flags and d.flags.get("is_verified", False)
+                }
+
+                # If user's domain is not in the verified domains, deny access
+                if email_domain not in verified_domain_slugs:
+                    return {
+                        "error": "AUTH_DOMAIN_DENIED",
+                        "message": f"Your email domain '{email_domain}' is not allowed for this organization",
+                        "current_domain": email_domain,
+                        "allowed_domains": list(verified_domain_slugs),
+                    }
+
         return None
 
     def _matches_policy(