tgex-backend/src/domain/workspace_permissions.py

from __future__ import annotations

from dataclasses import dataclass
from uuid import UUID

from . import WorkspaceAccessDenied
from .workspace import (
    PermissionKey,
    PermissionScopeType,
    Workspace,
    WorkspaceUser,
)


@dataclass
class WorkspacePermissions:
    global_permissions: set[PermissionKey]
    scoped_permissions: dict[tuple[PermissionKey, PermissionScopeType], set[UUID]]

    @classmethod
    def from_membership(cls, membership: WorkspaceUser) -> WorkspacePermissions:
        global_permissions = {permission.permission for permission in getattr(membership, 'permissions', []) or []}

        scoped_permissions: dict[tuple[PermissionKey, PermissionScopeType], set[UUID]] = {}
        for scope in getattr(membership, 'permission_scopes', []) or []:
            key = (scope.permission, scope.scope_type)
            scoped_permissions.setdefault(key, set()).add(scope.scope_id)

        return cls(global_permissions=global_permissions, scoped_permissions=scoped_permissions)

    def has_global(self, permission: PermissionKey) -> bool:
        return permission in self.global_permissions or PermissionKey.ADMIN_FULL in self.global_permissions

    def allowed_project_ids(self, permission: PermissionKey) -> set[UUID] | None:
        if self.has_global(permission):
            return None

        allowed: set[UUID] = set()
        for key in (permission, PermissionKey.ADMIN_FULL):
            project_scope = self.scoped_permissions.get((key, PermissionScopeType.PROJECT))
            if project_scope:
                allowed.update(project_scope)

        return allowed

    def allowed_creative_ids(self, permission: PermissionKey) -> set[UUID] | None:
        if self.has_global(permission):
            return None

        allowed: set[UUID] = set()
        for key in (permission, PermissionKey.ADMIN_FULL):
            creative_scope = self.scoped_permissions.get((key, PermissionScopeType.CREATIVE))
            if creative_scope:
                allowed.update(creative_scope)

        return allowed

    def allowed_placement_ids(self, permission: PermissionKey) -> set[UUID] | None:
        if self.has_global(permission):
            return None

        allowed: set[UUID] = set()
        for key in (permission, PermissionKey.ADMIN_FULL):
            placement_scope = self.scoped_permissions.get((key, PermissionScopeType.PLACEMENT))
            if placement_scope:
                allowed.update(placement_scope)

        return allowed

    def allowed_channel_ids(self, permission: PermissionKey) -> set[UUID] | None:
        if self.has_global(permission):
            return None

        allowed: set[UUID] = set()
        for key in (permission, PermissionKey.ADMIN_FULL):
            channel_scope = self.scoped_permissions.get((key, PermissionScopeType.CHANNEL))
            if channel_scope:
                allowed.update(channel_scope)

        return allowed

    def has_permission(
        self,
        permission: PermissionKey,
        *,
        scope_type: PermissionScopeType | None = None,
        scope_id: UUID | None = None,
    ) -> bool:
        if self.has_global(permission):
            return True

        if scope_type is not None and scope_id is not None:
            if scope_type == PermissionScopeType.PROJECT:
                allowed = self.allowed_project_ids(permission)
            elif scope_type == PermissionScopeType.CREATIVE:
                allowed = self.allowed_creative_ids(permission)
            elif scope_type == PermissionScopeType.PLACEMENT:
                allowed = self.allowed_placement_ids(permission)
            elif scope_type == PermissionScopeType.CHANNEL:
                allowed = self.allowed_channel_ids(permission)
            else:
                return False

            if allowed is None:
                return True
            return scope_id in allowed

        return False

    def has_any(self, permission: PermissionKey) -> bool:
        allowed = self.allowed_project_ids(permission)
        if allowed is None:
            return True
        return bool(allowed)


@dataclass
class WorkspacePermissionContext:
    workspace: Workspace
    membership: WorkspaceUser
    permissions: WorkspacePermissions

    def allowed_project_ids(self, permission: PermissionKey) -> set[UUID] | None:
        return self.permissions.allowed_project_ids(permission)

    def allowed_creative_ids(self, permission: PermissionKey) -> set[UUID] | None:
        return self.permissions.allowed_creative_ids(permission)

    def allowed_placement_ids(self, permission: PermissionKey) -> set[UUID] | None:
        return self.permissions.allowed_placement_ids(permission)

    def allowed_channel_ids(self, permission: PermissionKey) -> set[UUID] | None:
        return self.permissions.allowed_channel_ids(permission)

    def ensure_project_permission(self, permission: PermissionKey, project_id: UUID) -> None:
        if not self.permissions.has_permission(
            permission,
            scope_type=PermissionScopeType.PROJECT,
            scope_id=project_id,
        ):
            raise WorkspaceAccessDenied(self.workspace.id)

    def ensure_creative_permission(self, permission: PermissionKey, creative_id: UUID) -> None:
        if not self.permissions.has_permission(
            permission,
            scope_type=PermissionScopeType.CREATIVE,
            scope_id=creative_id,
        ):
            raise WorkspaceAccessDenied(self.workspace.id)

    def ensure_placement_permission(self, permission: PermissionKey, placement_id: UUID) -> None:
        if not self.permissions.has_permission(
            permission,
            scope_type=PermissionScopeType.PLACEMENT,
            scope_id=placement_id,
        ):
            raise WorkspaceAccessDenied(self.workspace.id)

    def ensure_channel_permission(self, permission: PermissionKey, channel_id: UUID) -> None:
        if not self.permissions.has_permission(
            permission,
            scope_type=PermissionScopeType.CHANNEL,
            scope_id=channel_id,
        ):
            raise WorkspaceAccessDenied(self.workspace.id)


def build_workspace_permission_context(membership: WorkspaceUser) -> WorkspacePermissionContext:
    if membership.workspace is None:
        raise ValueError('Workspace relation must be loaded for membership permissions')

    permissions = WorkspacePermissions.from_membership(membership)

    return WorkspacePermissionContext(
        workspace=membership.workspace,
        membership=membership,
        permissions=permissions,
    )