ITWorx_MCP_Hub-Public/services/control_api/persistence/machine_client.py

"""Idempotent provisioning for real Authentik-backed MCP machine clients."""

from __future__ import annotations

from dataclasses import dataclass
from uuid import NAMESPACE_URL, UUID, uuid5

from sqlalchemy import select
from sqlalchemy.orm import Session

from packages.hub_clients import ClientKind

from .models import (
    ClientBindingRow,
    ClientRow,
    EnvironmentRow,
    ProjectRow,
    VirtualServerRow,
    WorkspaceRow,
)


@dataclass(frozen=True)
class MachineClientResult:
    client_id: UUID
    binding_id: UUID
    endpoint: str


def provision_machine_client(
    session: Session,
    *,
    issuer: str,
    subject: str,
    oauth_username: str,
    name: str,
    project_slug: str,
    environment_name: str = "prod",
    workspace_slug: str = "itworx",
    public_host: str = "mcp.example.com",
    client_id: UUID | None = None,
) -> MachineClientResult:
    """Bind one validated Authentik subject to exactly one external virtual server.

    The caller obtains ``subject`` from a validated Authentik token or the exact
    Authentik user UUID. Credentials are deliberately not accepted here.
    """

    issuer = issuer.strip()
    subject = subject.strip()
    oauth_username = oauth_username.strip()
    name = name.strip()
    if not issuer.startswith("https://") or not subject or not oauth_username or not name:
        raise ValueError("machine client identity is incomplete")
    workspace = session.scalar(select(WorkspaceRow).where(WorkspaceRow.slug == workspace_slug))
    if workspace is None:
        raise ValueError("machine client workspace is not registered")
    project = session.scalar(
        select(ProjectRow).where(
            ProjectRow.workspace_id == workspace.id,
            ProjectRow.slug == project_slug,
            ProjectRow.status == "active",
        )
    )
    if project is None:
        raise ValueError("machine client project is not registered")
    environment = session.scalar(
        select(EnvironmentRow).where(
            EnvironmentRow.workspace_id == workspace.id,
            EnvironmentRow.project_id == project.id,
            EnvironmentRow.name == environment_name,
            EnvironmentRow.status == "active",
        )
    )
    if environment is None:
        raise ValueError("machine client environment is not registered")
    servers = session.scalars(
        select(VirtualServerRow).where(
            VirtualServerRow.workspace_id == workspace.id,
            VirtualServerRow.project_id == project.id,
            VirtualServerRow.environment_id == environment.id,
            VirtualServerRow.enabled.is_(True),
            VirtualServerRow.gateway_external_id.is_not(None),
        )
    ).all()
    if len(servers) != 1:
        raise ValueError(
            "machine client virtual server is not uniquely determined "
            f"({len(servers)} externally-routed servers found)"
        )
    server = servers[0]
    deterministic_id = uuid5(NAMESPACE_URL, f"{issuer}\n{subject}")
    principal = session.scalar(
        select(ClientRow).where(
            ClientRow.workspace_id == workspace.id,
            ClientRow.issuer == issuer,
            ClientRow.subject == subject,
        )
    )
    selected_id = client_id or deterministic_id
    by_id = session.get(ClientRow, selected_id)
    if principal is not None and principal.id != selected_id:
        raise ValueError("machine client principal is already managed by another identity")
    if (
        client_id is None
        and by_id is not None
        and (by_id.issuer != issuer or by_id.subject != subject)
    ):
        raise ValueError("machine client deterministic identity collision")
    if client_id is not None and by_id is not None:
        existing_metadata = by_id.record_metadata if isinstance(by_id.record_metadata, dict) else {}
        if (
            existing_metadata.get("managedBy") != "runtime-onboarding"
            or existing_metadata.get("project") != project_slug
            or existing_metadata.get("environment") != environment_name
            or by_id.status != "pending"
        ):
            raise ValueError("existing Hub client is not an activatable pending identity")
        try:
            kind = ClientKind(existing_metadata.get("kind", ClientKind.STDIO_BRIDGE.value))
        except (TypeError, ValueError) as exc:
            raise ValueError("pending machine client has an invalid client kind") from exc
    else:
        kind = ClientKind.STDIO_BRIDGE
    client = principal or by_id
    metadata = {
        "managedBy": "authentik-machine-provisioning",
        "identityProvisioned": True,
        "name": name,
        "project": project_slug,
        "environment": environment_name,
        "kind": kind.value,
        "oauthUsername": oauth_username,
        "credentialReference": f"authentik://service-account/{oauth_username}",
        "operatorDisabled": False,
    }
    if client is None:
        client = ClientRow(
            id=selected_id,
            workspace_id=workspace.id,
            subject=subject,
            issuer=issuer,
            client_type="mcp",
            status="active",
            record_metadata=metadata,
        )
        session.add(client)
        session.flush()
    else:
        if client.status == "revoked":
            raise ValueError("revoked machine client cannot be reprovisioned")
        client.issuer = issuer
        client.subject = subject
        client.status = "active"
        client.record_metadata = metadata
        session.flush()
    bindings = session.scalars(
        select(ClientBindingRow).where(ClientBindingRow.client_id == client.id)
    ).all()
    foreign = [item for item in bindings if item.virtual_server_id != server.id]
    if foreign:
        raise ValueError("machine client already has a different project binding")
    binding = next(iter(bindings), None)
    if binding is None:
        binding = ClientBindingRow(
            workspace_id=workspace.id,
            project_id=project.id,
            client_id=client.id,
            virtual_server_id=server.id,
            enabled=True,
        )
        session.add(binding)
    else:
        binding.enabled = True
    session.flush()
    endpoint = f"https://{public_host}/mcp/{project_slug}/{environment_name}"
    return MachineClientResult(client_id=client.id, binding_id=binding.id, endpoint=endpoint)