Coverage for paperless/consumers.py: 26%

1from __future__ import annotations 

2 

3import asyncio 

4import contextlib 

5import json 

6from typing import TYPE_CHECKING 

7 

8from channels.generic.websocket import AsyncWebsocketConsumer 

9 

10if TYPE_CHECKING: 10 ↛ 11line 10 didn't jump to line 11 because the condition on line 10 was never true

11 from django.contrib.auth.base_user import AbstractBaseUser 

12 from django.contrib.auth.models import AnonymousUser 

13 

14 from documents.plugins.helpers import DocumentsDeletedPayload 

15 from documents.plugins.helpers import DocumentUpdatedPayload 

16 from documents.plugins.helpers import PermissionsData 

17 from documents.plugins.helpers import StatusUpdatePayload 

18 

19HEARTBEAT_INTERVAL = 30 

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

21 

22 

23class StatusConsumer(AsyncWebsocketConsumer): 

24 heartbeat_task: asyncio.Task | None = None 

25 

26 def _authenticated(self) -> bool: 

27 user: AbstractBaseUser | AnonymousUser | None = self.scope.get("user") 

28 return user is not None and user.is_authenticated 

29 

30 async def _can_view(self, data: PermissionsData) -> bool: 

31 user: AbstractBaseUser | AnonymousUser | None = self.scope.get("user") 

32 if user is None: 

33 return False 

34 owner_id = data.get("owner_id") 

35 users_can_view = data.get("users_can_view", []) 

36 groups_can_view = data.get("groups_can_view", []) 

37 

38 if user.is_superuser or user.id == owner_id or user.id in users_can_view: 

39 return True 

40 

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

42 

43 async def connect(self) -> None: 

44 if not self._authenticated(): 

45 await self.close() 

46 return 

47 await self.channel_layer.group_add("status_updates", self.channel_name) 

48 await self.accept() 

49 self._start_heartbeat() 

50 

51 async def disconnect(self, code: int) -> None: 

52 await self._stop_heartbeat() 

53 await self.channel_layer.group_discard("status_updates", self.channel_name) 

54 

55 def _start_heartbeat(self) -> None: 

56 self.heartbeat_task = asyncio.create_task(self._heartbeat_loop()) 

57 

58 async def _stop_heartbeat(self) -> None: 

59 if self.heartbeat_task is not None: 

60 self.heartbeat_task.cancel() 

61 with contextlib.suppress(asyncio.CancelledError): 

62 await self.heartbeat_task 

63 self.heartbeat_task = None 

64 

65 async def _heartbeat_loop(self) -> None: 

66 while True: 

67 await asyncio.sleep(HEARTBEAT_INTERVAL) 

68 await self.send(HEARTBEAT_MESSAGE) 

69 

70 async def status_update(self, event: StatusUpdatePayload) -> None: 

71 if not self._authenticated(): 

72 await self.close() 

73 elif await self._can_view(event["data"]): 

74 await self.send(json.dumps(event)) 

75 

76 async def documents_deleted(self, event: DocumentsDeletedPayload) -> None: 

77 if not self._authenticated(): 

78 await self.close() 

79 else: 

80 await self.send(json.dumps(event)) 

81 

82 async def document_updated(self, event: DocumentUpdatedPayload) -> None: 

83 if not self._authenticated(): 

84 await self.close() 

85 elif await self._can_view(event["data"]): 

86 await self.send(json.dumps(event))