From 9c00c86e9bf709fc6aa2ff179b765123f4db060a Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Thu, 26 Nov 2020 17:08:26 +0100
Subject: [PATCH] root: add Channels Message Storage and consumer

---
 passbook/root/messages/__init__.py |  0
 passbook/root/messages/consumer.py | 20 ++++++++++++++++++
 passbook/root/messages/storage.py  | 34 ++++++++++++++++++++++++++++++
 passbook/root/settings.py          |  2 ++
 passbook/root/websocket.py         |  7 +++++-
 5 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 passbook/root/messages/__init__.py
 create mode 100644 passbook/root/messages/consumer.py
 create mode 100644 passbook/root/messages/storage.py

diff --git a/passbook/root/messages/__init__.py b/passbook/root/messages/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/passbook/root/messages/consumer.py b/passbook/root/messages/consumer.py
new file mode 100644
index 000000000..04874c1cf
--- /dev/null
+++ b/passbook/root/messages/consumer.py
@@ -0,0 +1,20 @@
+"""websocket Message consumer"""
+from channels.generic.websocket import JsonWebsocketConsumer
+from django.core.cache import cache
+
+
+class MessageConsumer(JsonWebsocketConsumer):
+    """Consumer which sends django.contrib.messages Messages over WS.
+    channel_name is saved into cache with user_id, and when a add_message is called"""
+
+    def connect(self):
+        self.accept()
+        cache.set(f"user_{self.scope['user'].pk}_{self.channel_name}", True)
+
+    # pylint: disable=unused-argument
+    def disconnect(self, close_code):
+        cache.delete(f"user_{self.scope['user'].pk}_{self.channel_name}")
+
+    def event_update(self, event: dict):
+        """Event handler which is called by Messages Storage backend"""
+        self.send_json(event)
diff --git a/passbook/root/messages/storage.py b/passbook/root/messages/storage.py
new file mode 100644
index 000000000..0b7a40b44
--- /dev/null
+++ b/passbook/root/messages/storage.py
@@ -0,0 +1,34 @@
+"""Channels Messages storage"""
+from asgiref.sync import async_to_sync
+from channels.layers import get_channel_layer
+from django.contrib.messages.storage.base import BaseStorage, Message
+from django.core.cache import cache
+from django.http.request import HttpRequest
+
+
+class ChannelsStorage(BaseStorage):
+    """Send contrib.messages over websocket"""
+
+    def __init__(self, request: HttpRequest) -> None:
+        # pyright: reportGeneralTypeIssues=false
+        super().__init__(request)
+        self.channel = get_channel_layer()
+
+    def _store(self, messages: list[Message], response, *args, **kwargs):
+        prefix = f"user_{self.request.user.pk}_"
+        keys = cache.keys(f"{prefix}*")
+        for key in keys:
+            uid = key.replace(prefix, "")
+            for message in messages:
+                async_to_sync(self.channel.send)(
+                    uid,
+                    {
+                        "type": "event.update",
+                        "levelTag": message.level_tag,
+                        "tags": message.tags,
+                        "message": message.message,
+                    },
+                )
+
+    def _get(self, *args, **kwargs):
+        return [], True
diff --git a/passbook/root/settings.py b/passbook/root/settings.py
index 0dff2c367..d454bf05c 100644
--- a/passbook/root/settings.py
+++ b/passbook/root/settings.py
@@ -177,6 +177,8 @@ SESSION_ENGINE = "django.contrib.sessions.backends.cache"
 SESSION_CACHE_ALIAS = "default"
 SESSION_COOKIE_SAMESITE = "lax"
 
+MESSAGE_STORAGE = "passbook.root.messages.storage.ChannelsStorage"
+
 MIDDLEWARE = [
     "django_prometheus.middleware.PrometheusBeforeMiddleware",
     "django.contrib.sessions.middleware.SessionMiddleware",
diff --git a/passbook/root/websocket.py b/passbook/root/websocket.py
index 005de97f1..31b518110 100644
--- a/passbook/root/websocket.py
+++ b/passbook/root/websocket.py
@@ -1,6 +1,11 @@
 """root Websocket URLS"""
+from channels.auth import AuthMiddlewareStack
 from django.urls import path
 
 from passbook.outposts.channels import OutpostConsumer
+from passbook.root.messages.consumer import MessageConsumer
 
-websocket_urlpatterns = [path("ws/outpost/<uuid:pk>/", OutpostConsumer.as_asgi())]
+websocket_urlpatterns = [
+    path("ws/outpost/<uuid:pk>/", OutpostConsumer.as_asgi()),
+    path("ws/client/", AuthMiddlewareStack(MessageConsumer.as_asgi())),
+]