robosats/chat/consumers.py

from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from api.models import Order
from chat.models import ChatRoom, Message

import json

class ChatRoomConsumer(AsyncWebsocketConsumer):

    @database_sync_to_async
    def allow_in_chatroom(self):
        order = Order.objects.get(id=self.order_id)

        if not order.status in [Order.Status.CHA, Order.Status.FSE]:
            print("Order is not in chat status")
            return False

        if not (order.maker == self.user or order.taker == self.user):
            print("Not allowed in this chat")
            return False
        return True

    @database_sync_to_async
    def save_connect_user(self):
        '''Creates or updates the ChatRoom object'''

        order = Order.objects.get(id=self.order_id)

        if order.maker == self.user:
            ChatRoom.objects.update_or_create(
                id=self.order_id, 
                order=order,  
                room_group_name=self.room_group_name,
                defaults={
                    "maker": self.user,
                    "maker_connected": True,
                    }
                )

        elif order.taker == self.user:
            ChatRoom.objects.update_or_create(
                id=self.order_id, 
                order=order,  
                room_group_name=self.room_group_name,
                defaults={
                    "taker": self.user,
                    "taker_connected": True,
                    }
                )

        return None

    @database_sync_to_async
    def save_new_PGP_message(self, PGP_message):
        '''Creates a Message object'''

        order = Order.objects.get(id=self.order_id)
        chatroom = ChatRoom.objects.get(order=order)

        try:
            last_message = Message.objects.filter(order=order).latest()
            index = last_message.index + 1
        except:
            index = 1

        sender = self.scope["user"]
        if order.taker == sender:
            receiver = order.maker
        elif order.maker == sender:
            receiver = order.taker

        msg_obj = Message.objects.create(
                order=order,
                chatroom=chatroom,
                index=index,
                sender=sender,
                receiver=receiver,
                PGP_message=PGP_message,
                )
        return msg_obj

    @database_sync_to_async
    def save_disconnect_user(self):
        '''Creates or updates the ChatRoom object'''
        
        order = Order.objects.get(id=self.order_id)
        if order.maker == self.user:
            ChatRoom.objects.update_or_create(
                id=self.order_id, 
                defaults={
                    "maker_connected": False
                    }
                )
        elif order.taker == self.user:
            ChatRoom.objects.update_or_create(
                id=self.order_id,
                defaults={
                    "taker_connected": False
                    }
                )
        return None

    @database_sync_to_async
    def is_peer_connected(self):
        '''Returns whether the consumer's peer is connected'''

        chatroom = ChatRoom.objects.get(id=self.order_id)

        if chatroom.maker == self.user:
            return chatroom.taker_connected

        if chatroom.taker == self.user:
            return chatroom.maker_connected

    @database_sync_to_async
    def get_peer_PGP_public_key(self):
        '''Returns peer PGP public key'''

        order = Order.objects.get(id=self.order_id)

        if order.maker == self.user:
            return order.taker.profile.public_key

        if order.taker == self.user:
            return order.maker.profile.public_key

    @database_sync_to_async
    def get_all_PGP_messages(self):
        '''Returns all PGP messages'''

        order = Order.objects.get(id=self.order_id)
        messages = Message.objects.filter(order=order)

        msgs = []
        for message in messages:
            msgs.append({
                "index": message.index,
                "time": str(message.created_at),
                "message": message.PGP_message,
                "nick": str(message.sender),
                })

        return msgs

    async def connect(self):
        self.order_id = self.scope["url_route"]["kwargs"]["order_id"]
        self.room_group_name = f"chat_order_{self.order_id}"
        self.user = self.scope["user"]
        self.user_nick = str(self.user)

        allowed = await self.allow_in_chatroom()

        if allowed:
            await self.save_connect_user()
            await self.channel_layer.group_add(self.room_group_name,
                                            self.channel_name)

            await self.accept()

            # Send peer PGP public keys
            peer_public_key = await self.get_peer_PGP_public_key()
            peer_connected = await self.is_peer_connected()
            await self.channel_layer.group_send(
                self.room_group_name,
                {
                    "type": "chatroom_message",
                    "message": peer_public_key,
                    "nick": self.scope["user"].username,
                    "peer_connected": peer_connected,
                },
            )

    async def disconnect(self, close_code):
        await self.save_disconnect_user()
        await self.channel_layer.group_discard(self.room_group_name,
                                               self.channel_name)
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                "type": "chatroom_message",
                "message": 'peer-disconnected',
                "nick": self.scope["user"].username,
                "peer_connected": False,
            },
        )

    async def receive(self, text_data):
        text_data_json = json.loads(text_data)
        message = text_data_json["message"]
        peer_connected = await self.is_peer_connected()
        
        # Encrypted messages are stored. They are served later when a user reconnects.
        if message[0:27] == '-----BEGIN PGP MESSAGE-----': 
            # save to database
            msg_obj = await self.save_new_PGP_message(message)

            index = msg_obj.index
            message = msg_obj.PGP_message
            time = str(msg_obj.created_at)

            await self.channel_layer.group_send(
                self.room_group_name,
                {
                    "type": "PGP_message",
                    "index": index,
                    "message": message,
                    "time": time,
                    "nick": self.scope["user"].username,
                    "peer_connected": peer_connected,
                },
            )
        
        # Encrypted messages are served when the user requests them
        elif message[0:23] == '-----SERVE HISTORY-----': 
            # If there is any stored message, serve them.
            msgs = await self.get_all_PGP_messages()
            peer_connected = await self.is_peer_connected()
            for msg in msgs:
                await self.channel_layer.group_send(
                    self.room_group_name,
                    {
                        "type": "PGP_message",
                        "index": msg['index'],
                        "time": msg['time'],
                        "message": msg['message'],
                        "nick": msg['nick'],
                        "peer_connected": peer_connected,
                    },
                )
        # Unencrypted messages are not stored, just echoed.
        else:
            await self.channel_layer.group_send(
                self.room_group_name,
                {
                    "type": "chatroom_message",
                    "message": message,
                    "nick": self.scope["user"].username,
                    "peer_connected": peer_connected,
                },
            )

    async def chatroom_message(self, event):
        message = event["message"]
        nick = event["nick"]
        peer_connected = event["peer_connected"]

        await self.send(text_data=json.dumps({
            "message": message,
            "user_nick": nick,
            "peer_connected": peer_connected,
        }))

    async def PGP_message(self, event):
        message = event["message"]
        nick = event["nick"]
        index = event["index"]
        peer_connected = event["peer_connected"]
        time = event["time"]

        await self.send(text_data=json.dumps({
            "index": index,
            "message": message,
            "user_nick": nick,
            "peer_connected": peer_connected,
            "time":time,
        }))