From 2a4e40b41b8471443b75ebcee69bfd16b53cbf88 Mon Sep 17 00:00:00 2001 From: BarsTiger Date: Sun, 2 Jul 2023 11:54:49 +0300 Subject: [PATCH] Working development server (no encryption now) --- dragonion_server/__init__.py | 5 +- dragonion_server/cli/__init__.py | 6 ++ dragonion_server/cli/cmd/__init__.py | 1 + dragonion_server/cli/cmd/service/__init__.py | 1 + dragonion_server/cli/cmd/service/remove.py | 37 ++++++++ dragonion_server/cli/cmd/service/run.py | 37 ++++++++ dragonion_server/cli/cmd/service/service.py | 14 +++ dragonion_server/cli/cmd/service/write.py | 37 ++++++++ dragonion_server/cli/common.py | 10 ++ dragonion_server/cli/utils/__init__.py | 6 ++ dragonion_server/cli/utils/groups.py | 18 ++++ dragonion_server/common.py | 3 + .../modules/server/handlers/__init__.py | 1 + .../server/handlers/managers/__init__.py | 1 + .../server/handlers/managers/connection.py | 36 +++++++ .../modules/server/handlers/managers/room.py | 95 +++++++++++++++++++ .../server/handlers/managers/service.py | 52 ++++++++++ .../server/handlers/objects/__init__.py | 1 + .../server/handlers/objects/webmessage.py | 69 ++++++++++++++ .../server/handlers/websocket_server.py | 44 +++++++++ dragonion_server/modules/server/routes.py | 10 +- dragonion_server/modules/server/server.py | 4 +- pyproject.toml | 7 +- 23 files changed, 486 insertions(+), 9 deletions(-) create mode 100644 dragonion_server/cli/__init__.py create mode 100644 dragonion_server/cli/cmd/__init__.py create mode 100644 dragonion_server/cli/cmd/service/__init__.py create mode 100644 dragonion_server/cli/cmd/service/remove.py create mode 100644 dragonion_server/cli/cmd/service/run.py create mode 100644 dragonion_server/cli/cmd/service/service.py create mode 100644 dragonion_server/cli/cmd/service/write.py create mode 100644 dragonion_server/cli/common.py create mode 100644 dragonion_server/cli/utils/__init__.py create mode 100644 dragonion_server/cli/utils/groups.py create mode 100644 dragonion_server/common.py create mode 100644 dragonion_server/modules/server/handlers/__init__.py create mode 100644 dragonion_server/modules/server/handlers/managers/__init__.py create mode 100644 dragonion_server/modules/server/handlers/managers/connection.py create mode 100644 dragonion_server/modules/server/handlers/managers/room.py create mode 100644 dragonion_server/modules/server/handlers/managers/service.py create mode 100644 dragonion_server/modules/server/handlers/objects/__init__.py create mode 100644 dragonion_server/modules/server/handlers/objects/webmessage.py create mode 100644 dragonion_server/modules/server/handlers/websocket_server.py diff --git a/dragonion_server/__init__.py b/dragonion_server/__init__.py index 00a1786..883b1f7 100644 --- a/dragonion_server/__init__.py +++ b/dragonion_server/__init__.py @@ -1,2 +1,5 @@ +from dragonion_server.cli import cli + + def main(): - print('DRAGONION-SERVER') + cli() diff --git a/dragonion_server/cli/__init__.py b/dragonion_server/cli/__init__.py new file mode 100644 index 0000000..e82a8ca --- /dev/null +++ b/dragonion_server/cli/__init__.py @@ -0,0 +1,6 @@ +from .common import cli + + +__all__ = [ + 'cli' +] diff --git a/dragonion_server/cli/cmd/__init__.py b/dragonion_server/cli/cmd/__init__.py new file mode 100644 index 0000000..2ae2839 --- /dev/null +++ b/dragonion_server/cli/cmd/__init__.py @@ -0,0 +1 @@ +pass diff --git a/dragonion_server/cli/cmd/service/__init__.py b/dragonion_server/cli/cmd/service/__init__.py new file mode 100644 index 0000000..2ae2839 --- /dev/null +++ b/dragonion_server/cli/cmd/service/__init__.py @@ -0,0 +1 @@ +pass diff --git a/dragonion_server/cli/cmd/service/remove.py b/dragonion_server/cli/cmd/service/remove.py new file mode 100644 index 0000000..31ef889 --- /dev/null +++ b/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) diff --git a/dragonion_server/cli/cmd/service/run.py b/dragonion_server/cli/cmd/service/run.py new file mode 100644 index 0000000..e44f238 --- /dev/null +++ b/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) diff --git a/dragonion_server/cli/cmd/service/service.py b/dragonion_server/cli/cmd/service/service.py new file mode 100644 index 0000000..75fc2ff --- /dev/null +++ b/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() + } +) diff --git a/dragonion_server/cli/cmd/service/write.py b/dragonion_server/cli/cmd/service/write.py new file mode 100644 index 0000000..5e70b3c --- /dev/null +++ b/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) diff --git a/dragonion_server/cli/common.py b/dragonion_server/cli/common.py new file mode 100644 index 0000000..55bc9e0 --- /dev/null +++ b/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() + ] +) diff --git a/dragonion_server/cli/utils/__init__.py b/dragonion_server/cli/utils/__init__.py new file mode 100644 index 0000000..226d155 --- /dev/null +++ b/dragonion_server/cli/utils/__init__.py @@ -0,0 +1,6 @@ +from .groups import ModuleGroup + + +__all__ = [ + 'ModuleGroup' +] diff --git a/dragonion_server/cli/utils/groups.py b/dragonion_server/cli/utils/groups.py new file mode 100644 index 0000000..08324d7 --- /dev/null +++ b/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) diff --git a/dragonion_server/common.py b/dragonion_server/common.py new file mode 100644 index 0000000..a9463af --- /dev/null +++ b/dragonion_server/common.py @@ -0,0 +1,3 @@ +from rich.console import Console + +console = Console() diff --git a/dragonion_server/modules/server/handlers/__init__.py b/dragonion_server/modules/server/handlers/__init__.py new file mode 100644 index 0000000..2ae2839 --- /dev/null +++ b/dragonion_server/modules/server/handlers/__init__.py @@ -0,0 +1 @@ +pass diff --git a/dragonion_server/modules/server/handlers/managers/__init__.py b/dragonion_server/modules/server/handlers/managers/__init__.py new file mode 100644 index 0000000..2ae2839 --- /dev/null +++ b/dragonion_server/modules/server/handlers/managers/__init__.py @@ -0,0 +1 @@ +pass diff --git a/dragonion_server/modules/server/handlers/managers/connection.py b/dragonion_server/modules/server/handlers/managers/connection.py new file mode 100644 index 0000000..b2ec10a --- /dev/null +++ b/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 + ) + ) diff --git a/dragonion_server/modules/server/handlers/managers/room.py b/dragonion_server/modules/server/handlers/managers/room.py new file mode 100644 index 0000000..41a4e42 --- /dev/null +++ b/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 diff --git a/dragonion_server/modules/server/handlers/managers/service.py b/dragonion_server/modules/server/handlers/managers/service.py new file mode 100644 index 0000000..6b57f36 --- /dev/null +++ b/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}' + ) diff --git a/dragonion_server/modules/server/handlers/objects/__init__.py b/dragonion_server/modules/server/handlers/objects/__init__.py new file mode 100644 index 0000000..2ae2839 --- /dev/null +++ b/dragonion_server/modules/server/handlers/objects/__init__.py @@ -0,0 +1 @@ +pass diff --git a/dragonion_server/modules/server/handlers/objects/webmessage.py b/dragonion_server/modules/server/handlers/objects/webmessage.py new file mode 100644 index 0000000..5da0abf --- /dev/null +++ b/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) diff --git a/dragonion_server/modules/server/handlers/websocket_server.py b/dragonion_server/modules/server/handlers/websocket_server.py new file mode 100644 index 0000000..aecb704 --- /dev/null +++ b/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 diff --git a/dragonion_server/modules/server/routes.py b/dragonion_server/modules/server/routes.py index a84be04..59e37b3 100644 --- a/dragonion_server/modules/server/routes.py +++ b/dragonion_server/modules/server/routes.py @@ -1,10 +1,10 @@ -from fastapi import APIRouter, Request -from fastapi.responses import PlainTextResponse +from fastapi import APIRouter, WebSocket +from .handlers.websocket_server import serve_websocket router = APIRouter() -@router.get("/", response_model=str) -async def root(request: Request): - return PlainTextResponse("dragonion-server") +@router.websocket("/{room_name}") +async def root(websocket: WebSocket, room_name: str): + await serve_websocket(websocket, room_name) diff --git a/dragonion_server/modules/server/server.py b/dragonion_server/modules/server/server.py index 71793b5..83da1aa 100644 --- a/dragonion_server/modules/server/server.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index dc14710..ffb4ce7 100644 --- a/pyproject.toml +++ b/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.98.0" +fastapi = "^0.99.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"