Working development server (no encryption now)

dragonion_server/__init__.py

@@ -1,2 +1,5 @@
+from dragonion_server.cli import cli
+
+
 def main():
-    print('DRAGONION-SERVER')
+    cli()

dragonion_server/cli/__init__.py

@@ -0,0 +1,6 @@
+from .common import cli
+
+
+__all__ = [
+    'cli'
+]

dragonion_server/cli/cmd/__init__.py

@@ -0,0 +1 @@
+pass

dragonion_server/cli/cmd/service/__init__.py

@@ -0,0 +1 @@
+pass

dragonion_server/cli/cmd/service/remove.py

@@ -0,0 +1,37 @@
+import os
+
+import click
+
+from dragonion_server.utils.config import db
+from dragonion_server.common import console
+
+
+class ServiceRemoveCommand(click.Command):
+    def __init__(self):
+        super().__init__(
+            name='remove',
+            callback=self.callback,
+            params=[
+                click.Option(
+                    ('--name', '-n'),
+                    required=True,
+                    prompt=True,
+                    type=str,
+                    help='Name of service to write to'
+                )
+            ]
+        )
+
+    @staticmethod
+    def callback(name: str):
+        try:
+            del db.services[name]
+            if os.path.isfile(f'{name}.auth'):
+                os.remove(f'{name}.auth')
+                
+            print(f'Removed service {name}')
+        except KeyError:
+            print(f'Service "{name}" does not exist in this storage')
+        except Exception as e:
+            assert e
+            console.print_exception(show_locals=True)

dragonion_server/cli/cmd/service/run.py

@@ -0,0 +1,37 @@
+import click
+
+from dragonion_server.modules.server import run
+from dragonion_server.common import console
+
+
+class ServiceRunCommand(click.Command):
+    def __init__(self):
+        super().__init__(
+            name='run',
+            callback=self.callback,
+            params=[
+                click.Option(
+                    ('--name', '-n'),
+                    required=True,
+                    prompt=True,
+                    type=str,
+                    help='Name of service to write to'
+                ),
+                click.Option(
+                    ('--port', '-p'),
+                    required=False,
+                    prompt=True,
+                    prompt_required=False,
+                    type=int,
+                    help='Port to start service on'
+                )
+            ]
+        )
+
+    @staticmethod
+    def callback(name: str, port: int | None):
+        try:
+            run(name, port)
+        except Exception as e:
+            assert e
+            console.print_exception(show_locals=True)

dragonion_server/cli/cmd/service/service.py

@@ -0,0 +1,14 @@
+from ...utils import ModuleGroup
+from .write import ServiceWriteCommand
+from .run import ServiceRunCommand
+from .remove import ServiceRemoveCommand
+
+
+service_group = ModuleGroup(
+    name='service',
+    commands={
+        'write': ServiceWriteCommand(),
+        'run': ServiceRunCommand(),
+        'remove': ServiceRemoveCommand()
+    }
+)

dragonion_server/cli/cmd/service/write.py

@@ -0,0 +1,37 @@
+import click
+
+from dragonion_server.utils.onion import Onion
+from dragonion_server.common import console
+
+
+class ServiceWriteCommand(click.Command):
+    def __init__(self):
+        super().__init__(
+            name='write',
+            callback=self.callback,
+            params=[
+                click.Option(
+                    ('--name', '-n'),
+                    required=True,
+                    prompt=True,
+                    type=str,
+                    help='Name of service to write to'
+                ),
+                click.Option(
+                    ('--port', '-p'),
+                    required=True,
+                    prompt=True,
+                    type=int,
+                    help='Port to start service on'
+                )
+            ]
+        )
+
+    @staticmethod
+    def callback(name: str, port: int):
+        try:
+            Onion.write_onion_service(name, port)
+            print(f'Written service "{name}" info')
+        except Exception as e:
+            assert e
+            console.print_exception(show_locals=True)

dragonion_server/cli/common.py

@@ -0,0 +1,10 @@
+import click
+from .cmd.service.service import service_group
+
+
+cli = click.CommandCollection(
+    name='dragonion-server',
+    sources=[
+        service_group()
+    ]
+)

dragonion_server/cli/utils/__init__.py

@@ -0,0 +1,6 @@
+from .groups import ModuleGroup
+
+
+__all__ = [
+    'ModuleGroup'
+]

dragonion_server/cli/utils/groups.py

@@ -0,0 +1,18 @@
+import click
+import typing as t
+
+
+class ModuleGroup(click.Group):
+    def __init__(
+        self,
+        name: t.Optional[str] = None,
+        commands: t.Optional[
+            t.Union[t.Dict[str, click.Command], t.Sequence[click.Command]]
+        ] = None,
+        **attrs: t.Any,
+    ) -> None:
+        new_commands = dict()
+        for command_key in commands.keys():
+            new_commands[f'{name}-{command_key}'] = commands[command_key]
+            
+        super().__init__(name, new_commands, **attrs)

dragonion_server/common.py

@@ -0,0 +1,3 @@
+from rich.console import Console
+
+console = Console()

dragonion_server/modules/server/handlers/__init__.py

@@ -0,0 +1 @@
+pass

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

@@ -0,0 +1 @@
+pass

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

@@ -0,0 +1,36 @@
+from attrs import define
+from fastapi import WebSocket
+from ..objects.webmessage import (
+    webmessages_union,
+    webmessage_error_message_literal,
+    WebErrorMessage,
+    WebUserMessage
+)
+
+
+@define
+class Connection(object):
+    ws: WebSocket
+    username: str
+    public_key: str
+
+    async def send_webmessage(self, obj: webmessages_union):
+        await self.ws.send_text(obj.to_json())
+
+    async def send_error(
+            self,
+            error_message: webmessage_error_message_literal
+    ):
+        await self.send_webmessage(
+            WebErrorMessage(
+                error_message=error_message
+            )
+        )
+
+    async def send_connect(self):
+        await self.send_webmessage(
+            WebUserMessage(
+                type="connect",
+                username=self.username
+            )
+        )

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

@@ -0,0 +1,95 @@
+from attrs import define
+from .connection import Connection
+from typing import Dict
+from fastapi import WebSocket
+
+from ..objects.webmessage import (
+    webmessages_union,
+    WebMessageMessage,
+    WebNotificationMessage,
+    webmessage_error_message_literal,
+    WebErrorMessage, 
+    WebUserMessage
+)
+
+
+@define
+class Room(object):
+    connections: Dict[str, Connection] = {}
+
+    async def accept_connection(self, ws: WebSocket) -> Connection:
+        print('Incoming connection')
+        await ws.accept()
+        connection = Connection(
+            username=(username := await ws.receive_text()),
+            ws=ws,
+            public_key=''
+        )
+        if username in self.connections.keys():
+            await connection.send_error(
+                'username_exists'
+            )
+
+        self.connections[username] = connection
+        await connection.send_connect()
+        print(f'Accepted {username}')
+        return connection
+
+    async def broadcast_webmessage(self, obj: webmessages_union):
+        for connection in self.connections.values():
+            print(f'Sending to {connection.username}: {obj}')
+            await connection.send_webmessage(obj)
+
+    async def broadcast_message(self, from_username: str, message: str):
+        await self.broadcast_webmessage(
+            WebMessageMessage(
+                username=from_username,
+                message=message
+            )
+        )
+
+    async def broadcast_notification(self, message: str):
+        await self.broadcast_webmessage(
+            WebNotificationMessage(
+                message=message
+            )
+        )
+
+    async def broadcast_error(
+            self,
+            error_message: webmessage_error_message_literal
+    ):
+        await self.broadcast_webmessage(
+            WebErrorMessage(
+                error_message=error_message
+            )
+        )
+
+    async def broadcast_user_disconnected(self, username: str):
+        await self.broadcast_webmessage(
+            WebUserMessage(
+                type="disconnect",
+                username=username
+            )
+        )
+
+    async def get_connection_by(self, attribute: str, value: str) -> Connection | None:
+        for connection in self.connections.values():
+            if getattr(connection, attribute) == value:
+                return connection
+
+    async def disconnect(self, connection: Connection, close_reason: str | None = None):
+        if connection not in self.connections.values():
+            return
+
+        del self.connections[connection.username]
+
+        try:
+            await connection.ws.close(
+                reason=close_reason
+            )
+        except Exception as e:
+            print(f'Cannot close {connection.username} ws, '
+                  f'maybe it\'s already closed: {e}')
+
+        return connection.username

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

@@ -0,0 +1,52 @@
+from .connection import Connection
+from .room import Room
+from typing import Dict
+
+from ..objects.webmessage import (
+    webmessage_error_message_literal
+)
+
+
+class Service(object):
+    rooms: Dict[str, Room] = {}
+
+    async def get_room(self, name: str) -> Room:
+        if name in self.rooms.keys():
+            return self.rooms[name]
+
+        self.rooms[name] = Room()
+        return self.rooms[name]
+
+    async def broadcast_notification(self, message: str):
+        for room in self.rooms.values():
+            await room.broadcast_notification(message)
+
+    async def broadcast_error(
+            self,
+            error_message: webmessage_error_message_literal
+    ):
+        for room in self.rooms.values():
+            await room.broadcast_error(error_message)
+
+    async def get_room_by_connection(self, connection: Connection) -> Room:
+        for room in self.rooms.values():
+            if connection in room.connections.values():
+                return room
+
+    async def get_connection_by_attribute(
+            self, attribute: str, value: str
+    ) -> Connection:
+        for room in self.rooms.values():
+            if connection := await room.get_connection_by(attribute, value):
+                return connection
+
+    async def close_room(self, room_name: str, reason: str = 'Unknown reason'):
+        room = self.rooms.get(room_name)
+        if room is None:
+            return
+
+        for connection in room.connections.values():
+            await room.disconnect(
+                connection=connection,
+                close_reason=f'Room is closed: {reason}'
+            )

dragonion_server/modules/server/handlers/objects/__init__.py

@@ -0,0 +1 @@
+pass

dragonion_server/modules/server/handlers/objects/webmessage.py

@@ -0,0 +1,69 @@
+from typing import Literal, Final, Union
+from dataclasses_json import dataclass_json
+from dataclasses import dataclass
+
+
+webmessage_type_literal = Literal[
+    "connect", "message", "disconnect", "error", "notification"
+]
+webmessage_error_message_literal = Literal[
+    "unknown", "username_exists", "invalid_webmessage"
+]
+
+
+@dataclass_json
+@dataclass
+class _WebAnyMessage:
+    username: str | None = None
+    type: webmessage_type_literal = "message"
+    message: str | None = None
+    error_message: webmessage_error_message_literal | None = None
+
+
+@dataclass_json
+@dataclass
+class WebMessageMessage:
+    username: str
+    message: str
+    type: Final = "message"
+
+
+@dataclass_json
+@dataclass
+class WebErrorMessage:
+    error_message: webmessage_error_message_literal
+    type: Final = "error"
+
+
+@dataclass_json
+@dataclass
+class WebUserMessage:
+    type: Literal["connect", "disconnect"]
+    username: str
+
+
+@dataclass_json
+@dataclass
+class WebNotificationMessage:
+    message: str
+    type: Final = "notification"
+
+
+webmessages_union = Union[
+    WebMessageMessage,
+    WebErrorMessage,
+    WebUserMessage,
+    WebNotificationMessage
+]
+
+
+class WebMessage:
+    @staticmethod
+    def from_json(data) -> webmessages_union:
+        return {
+            "connect": WebUserMessage.from_json,
+            "disconnect": WebUserMessage.from_json,
+            "message": WebMessageMessage.from_json,
+            "error": WebErrorMessage.from_json,
+            "notification": WebNotificationMessage.from_json
+        }[_WebAnyMessage.from_json(data).type](data)

dragonion_server/modules/server/handlers/websocket_server.py

@@ -0,0 +1,44 @@
+from fastapi import WebSocket, WebSocketDisconnect
+from .managers.service import Service
+from .objects.webmessage import (
+    webmessages_union,
+    WebMessage
+)
+
+
+service = Service()
+
+
+async def serve_websocket(websocket: WebSocket, room_name: str):
+    print(f'Connection opened room {room_name}')
+    room = await service.get_room(room_name)
+    connection = await room.accept_connection(websocket)
+
+    while True:
+        try:
+            data = await websocket.receive_text()
+            print(f"Received in {room_name}: ", data)
+
+            try:
+                webmessage: webmessages_union = \
+                    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)
+
+        except RuntimeError:
+            username = await room.disconnect(connection)
+            await room.broadcast_user_disconnected(username)
+            print(f'Closed {username}')
+            break
+        except WebSocketDisconnect:
+            username = await room.disconnect(connection)
+            await room.broadcast_user_disconnected(username)
+            print(f'Closed {username}')
+            break
+        except Exception as e:
+            print(f'Exception in {connection.username}: {e.__class__}: {e}')
+            continue

dragonion_server/modules/server/routes.py

@@ -1,10 +1,10 @@
-from fastapi import APIRouter, Request
+from fastapi import APIRouter, WebSocket
-from fastapi.responses import PlainTextResponse
+from .handlers.websocket_server import serve_websocket
 
 
 router = APIRouter()
 
 
-@router.get("/", response_model=str)
+@router.websocket("/{room_name}")
-async def root(request: Request):
+async def root(websocket: WebSocket, room_name: str):
-    return PlainTextResponse("dragonion-server")
+    await serve_websocket(websocket, room_name)

dragonion_server/modules/server/server.py

@@ -14,7 +14,9 @@ def get_app(port: int, name: str) -> FastAPI:
     )
 
 
-def run(name: str, port: int = get_available_port()):
+def run(name: str, port: int | None = get_available_port()):
+    if port is None:
+        port = get_available_port()
     app = get_app(port, name)
     app.include_router(router)
     uvicorn.run(app, host='0.0.0.0', port=port)

pyproject.toml

@@ -12,10 +12,13 @@ psutil = "^5.9.5"
 pynacl = "^1.5.0"
 requests = "^2.31.0"
 sqlitedict = "^2.1.0"
-cleo = "^2.0.1"
+fastapi = "^0.99.0"
-fastapi = "^0.98.0"
 uvicorn = "^0.22.0"
 rich = "^13.4.2"
+click = "^8.1.3"
+attrs = "^23.1.0"
+dataclasses-json = "^0.5.9"
+websockets = "^11.0.3"
 
 [tool.poetry.scripts]
 dragonion-server = "dragonion_server:main"