paperless-ngx/src/paperless/consumers.py

import json

from asgiref.sync import async_to_sync
from channels.exceptions import AcceptConnection
from channels.exceptions import DenyConnection
from channels.generic.websocket import WebsocketConsumer


class StatusConsumer(WebsocketConsumer):
    def _authenticated(self):
        return "user" in self.scope and self.scope["user"].is_authenticated

    def _can_view(self, data):
        user = self.scope.get("user") if self.scope.get("user") else None
        owner_id = data.get("owner_id")
        users_can_view = data.get("users_can_view", [])
        groups_can_view = data.get("groups_can_view", [])
        return (
            user.is_superuser
            or user.id == owner_id
            or user.id in users_can_view
            or any(
                user.groups.filter(pk=group_id).exists() for group_id in groups_can_view
            )
        )

    def connect(self):
        if not self._authenticated():
            raise DenyConnection
        else:
            async_to_sync(self.channel_layer.group_add)(
                "status_updates",
                self.channel_name,
            )
            raise AcceptConnection

    def disconnect(self, close_code):
        async_to_sync(self.channel_layer.group_discard)(
            "status_updates",
            self.channel_name,
        )

    def status_update(self, event):
        if not self._authenticated():
            self.close()
        else:
            if self._can_view(event["data"]):
                self.send(json.dumps(event))

    def documents_deleted(self, event):
        if not self._authenticated():
            self.close()
        else:
            self.send(json.dumps(event))