paperless-ngx/src/paperless/consumers.py

from __future__ import annotations

import asyncio
import contextlib
import json
from typing import TYPE_CHECKING

from channels.generic.websocket import AsyncWebsocketConsumer

if TYPE_CHECKING:
    from django.contrib.auth.base_user import AbstractBaseUser
    from django.contrib.auth.models import AnonymousUser

    from documents.plugins.helpers import DocumentsDeletedPayload
    from documents.plugins.helpers import DocumentUpdatedPayload
    from documents.plugins.helpers import PermissionsData
    from documents.plugins.helpers import StatusUpdatePayload

HEARTBEAT_INTERVAL = 30
HEARTBEAT_MESSAGE = json.dumps({"type": "heartbeat"})


class StatusConsumer(AsyncWebsocketConsumer):
    heartbeat_task: asyncio.Task | None = None

    def _authenticated(self) -> bool:
        user: AbstractBaseUser | AnonymousUser | None = self.scope.get("user")
        return user is not None and user.is_authenticated

    async def _can_view(self, data: PermissionsData) -> bool:
        user: AbstractBaseUser | AnonymousUser | None = self.scope.get("user")
        if user is None:
            return False
        owner_id = data.get("owner_id")
        users_can_view = data.get("users_can_view", [])
        groups_can_view = data.get("groups_can_view", [])

        if user.is_superuser or user.id == owner_id or user.id in users_can_view:
            return True

        return await user.groups.filter(pk__in=groups_can_view).aexists()

    async def connect(self) -> None:
        if not self._authenticated():
            await self.close()
            return
        await self.channel_layer.group_add("status_updates", self.channel_name)
        await self.accept()
        self._start_heartbeat()

    async def disconnect(self, code: int) -> None:
        await self._stop_heartbeat()
        await self.channel_layer.group_discard("status_updates", self.channel_name)

    def _start_heartbeat(self) -> None:
        self.heartbeat_task = asyncio.create_task(self._heartbeat_loop())

    async def _stop_heartbeat(self) -> None:
        if self.heartbeat_task is not None:
            self.heartbeat_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self.heartbeat_task
            self.heartbeat_task = None

    async def _heartbeat_loop(self) -> None:
        while True:
            await asyncio.sleep(HEARTBEAT_INTERVAL)
            await self.send(HEARTBEAT_MESSAGE)

    async def status_update(self, event: StatusUpdatePayload) -> None:
        if not self._authenticated():
            await self.close()
        elif await self._can_view(event["data"]):
            await self.send(json.dumps(event))

    async def documents_deleted(self, event: DocumentsDeletedPayload) -> None:
        if not self._authenticated():
            await self.close()
        else:
            await self.send(json.dumps(event))

    async def document_updated(self, event: DocumentUpdatedPayload) -> None:
        if not self._authenticated():
            await self.close()
        elif await self._can_view(event["data"]):
            await self.send(json.dumps(event))