Developing server, added run options, broadcasting encrypted works

dragonion_server/cli/cmd/service/run.py

@@ -1,6 +1,7 @@
+import sys
 import click
 
-from dragonion_server.modules.server import run
+from dragonion_server.modules.server import run, run_without_onion, integrate_onion
 from dragonion_server.common import console
 
 
@@ -24,14 +25,39 @@ class ServiceRunCommand(click.Command):
                     prompt_required=False,
                     type=int,
                     help='Port to start service on'
+                ),
+                click.Option(
+                    ('--without-tor', '-wt'),
+                    is_flag=True,
+                    help='Run service without tor'
+                ),
+                click.Option(
+                    ('--only-tor', '-ot'),
+                    is_flag=True,
+                    help='Run only tor proxy to service'
                 )
             ]
         )
 
     @staticmethod
-    def callback(name: str, port: int | None):
+    def callback(name: str, port: int | None, without_tor: bool, only_tor: bool):
         try:
-            run(name, port)
+            if without_tor and only_tor:
+                print('Cannot run only tor without tor, exiting')
+                sys.exit(1)
+            elif without_tor:
+                run_without_onion(name, port)
+            elif only_tor:
+                if port is None:
+                    print('For this mode, you need to specify port, '
+                          'to which requests will be redirected. Cannot start '
+                          'tor service, exiting')
+                    sys.exit(1)
+                onion = integrate_onion(port, name)
+                input('Press Enter to stop onion and service...')
+                onion.cleanup()
+            else:
+                run(name, port)
         except Exception as e:
             assert e
             console.print_exception(show_locals=True)

dragonion_server/modules/server/__init__.py

@@ -1,6 +1,9 @@
-from .server import run
+from .server import run, run_without_onion
+from .integration import integrate_onion
 
 
 __all__ = [
-    'run'
+    'run',
+    'run_without_onion',
+    'integrate_onion'
 ]

dragonion_server/modules/server/handlers/managers/connection.py

@@ -1,10 +1,9 @@
 from attrs import define
 from fastapi import WebSocket
-from dragonion_core.proto.web import (
+from dragonion_core.proto.web.webmessage import (
     webmessages_union,
     webmessage_error_message_literal,
-    WebErrorMessage,
+    WebErrorMessage
-    WebUserMessage
 )
 
 
@@ -36,15 +35,3 @@ class Connection(object):
                 error_message=error_message
             )
         )
-
-    async def send_connect(self):
-        """
-        When new user is connected, send info about user
-        :return: 
-        """
-        await self.send_webmessage(
-            WebUserMessage(
-                type="connect",
-                username=self.username
-            )
-        )

dragonion_server/modules/server/handlers/managers/exceptions.py

@@ -0,0 +1,2 @@
+class GotInvalidWebmessage(Exception):
+    pass

dragonion_server/modules/server/handlers/managers/room.py

@@ -1,15 +1,21 @@
 from attrs import define
 from .connection import Connection
+from .exceptions import GotInvalidWebmessage
 from typing import Dict
 from fastapi import WebSocket
 
-from dragonion_core.proto.web import (
+from json.decoder import JSONDecodeError
+
+from dragonion_core.proto.web.webmessage import (
     webmessages_union,
     WebMessageMessage,
+    WebBroadcastableMessage,
     WebNotificationMessage,
     webmessage_error_message_literal,
     WebErrorMessage,
-    WebUserMessage
+    WebConnectionMessage,
+    WebDisconnectMessage,
+    WebConnectionResultMessage
 )
 
 
@@ -17,7 +23,7 @@ from dragonion_core.proto.web import (
 class Room(object):
     connections: Dict[str, Connection] = {}
 
-    async def accept_connection(self, ws: WebSocket) -> Connection:
+    async def accept_connection(self, ws: WebSocket) -> Connection | None:
         """
         Accepts connection, checks username availability and adds it to dict of 
         connections 
@@ -26,19 +32,44 @@ class Room(object):
         """
         print('Incoming connection')
         await ws.accept()
+        try:
+            connection_message = WebConnectionMessage.from_json(
+                await ws.receive_text()
+            )
+        except JSONDecodeError:
+            await ws.send_text(WebErrorMessage(
+                'invalid_webmessage'
+            ).to_json())
+            await ws.close(reason='invalid_webmessage')
+            return 
+        
         connection = Connection(
-            username=(username := await ws.receive_text()),
+            username=connection_message.username,
             ws=ws,
-            public_key=''
+            public_key=connection_message.public_key
         )
-        if username in self.connections.keys():
+        
+        if connection_message.username in self.connections.keys():
             await connection.send_error(
                 'username_exists'
             )
+            await ws.close(reason='username_exists')
+            return 
 
-        self.connections[username] = connection
+        self.connections[connection_message.username] = connection
-        await connection.send_connect()
+        await connection.send_webmessage(WebConnectionResultMessage(
-        print(f'Accepted {username}')
+            connected_users=dict(
+                map(
+                    lambda i, j: (i, j),
+                    list(self.connections.keys()),
+                    [_connection.public_key for _connection 
+                     in self.connections.values()]
+                )
+            )
+        ))
+
+        await self.broadcast_webmessage(connection_message)
+        print(f'Accepted {connection_message.username}')
         return connection
 
     async def broadcast_webmessage(self, obj: webmessages_union):
@@ -51,19 +82,23 @@ class Room(object):
             print(f'Sending to {connection.username}: {obj}')
             await connection.send_webmessage(obj)
 
-    async def broadcast_message(self, from_username: str, message: str):
+    async def broadcast_message(self, broadcastable: WebBroadcastableMessage):
         """
         Broadcasts message to every user in room
-        :param from_username: User that sent message
+        :param broadcastable: String object with json representation of 
-        :param message: content
+                                   WebBroadcastableMessage
         :return: 
         """
-        await self.broadcast_webmessage(
+        try: 
-            WebMessageMessage(
+            for to_username in broadcastable.messages.keys():
-                username=from_username,
+                try:
-                message=message
+                    await self.connections[to_username].send_webmessage(
-            )
+                        broadcastable.messages[to_username]
-        )
+                    )
+                except KeyError:
+                    continue
+        except JSONDecodeError:
+            raise GotInvalidWebmessage
 
     async def broadcast_notification(self, message: str):
         """
@@ -99,8 +134,7 @@ class Room(object):
         :return:
         """
         await self.broadcast_webmessage(
-            WebUserMessage(
+            WebDisconnectMessage(
-                type="disconnect",
                 username=username
             )
         )

dragonion_server/modules/server/handlers/managers/service.py

@@ -2,7 +2,7 @@ from .connection import Connection
 from .room import Room
 from typing import Dict
 
-from dragonion_core.proto.web import (
+from dragonion_core.proto.web.webmessage import (
     webmessage_error_message_literal
 )
 

dragonion_server/modules/server/handlers/websocket_server.py

@@ -1,9 +1,10 @@
 from fastapi import WebSocket, WebSocketDisconnect
 from .managers.service import Service
-from dragonion_core.proto.web import (
+from dragonion_core.proto.web.webmessage import (
     webmessages_union,
     WebMessage
 )
+from .managers.exceptions import GotInvalidWebmessage
 
 
 service = Service()
@@ -23,26 +24,35 @@ async def serve_websocket(websocket: WebSocket, room_name: str):
     while True:
         try:
             data = await websocket.receive_text()
-            print(f"Received in {room_name}: ", data)
 
             try:
-                webmessage: webmessages_union = \
+                webmessage: webmessages_union = WebMessage.from_json(data)
-                    WebMessage.from_json(data)
             except Exception as e:
                 print(f"Cannot decode message, {e}")
                 await connection.send_error("invalid_webmessage")
                 continue
 
-            await room.broadcast_webmessage(webmessage)
+            try:
+                match webmessage.type:
+                    case "disconnect":
+                        await room.disconnect(connection)
+                    case "broadcastable":
+                        await room.broadcast_message(webmessage)
+
+            except GotInvalidWebmessage:
+                print('Invalid webmsg')
+                await connection.send_error("invalid_webmessage")
 
         except RuntimeError:
             username = await room.disconnect(connection)
-            await room.broadcast_user_disconnected(username)
+            if username is not None:
+                await room.broadcast_user_disconnected(username)
             print(f'Closed {username}')
             break
         except WebSocketDisconnect:
             username = await room.disconnect(connection)
-            await room.broadcast_user_disconnected(username)
+            if username is not None:
+                await room.broadcast_user_disconnected(username)
             print(f'Closed {username}')
             break
         except Exception as e:

dragonion_server/modules/server/server.py

@@ -32,3 +32,15 @@ def run(name: str, port: int | None = get_available_port()):
     app = get_app(port, name)
     app.include_router(router)
     uvicorn.run(app, host='0.0.0.0', port=port)
+
+
+def run_without_onion(name: str, port: int | None = get_available_port()):
+    if port is None:
+        port = get_available_port()
+        
+    app = FastAPI(
+        title=f'dragonion-server: {name}',
+        description=f'Secure dragonion chat endpoint server - service {name}'
+    )
+    app.include_router(router)
+    uvicorn.run(app, host='0.0.0.0', port=port)

pyproject.toml

@@ -6,7 +6,7 @@ authors = ["BarsTiger <zxcbarstiger@gmail.com>"]
 readme = "README.md"
 
 [tool.poetry.dependencies]
-python = "^3.10"
+python = ">=3.10,<3.12"
 stem = "^1.8.2"
 psutil = "^5.9.5"
 pynacl = "^1.5.0"