From ebbcf83cece91bbf5e6332db1778ffd2e3b230a4 Mon Sep 17 00:00:00 2001 From: BarsTiger Date: Sun, 23 Jul 2023 23:42:40 +0300 Subject: [PATCH] Developing server, added run options, broadcasting encrypted works --- dragonion_server/cli/cmd/service/run.py | 32 +++++++- dragonion_server/modules/server/__init__.py | 7 +- .../server/handlers/managers/connection.py | 17 +--- .../server/handlers/managers/exceptions.py | 2 + .../modules/server/handlers/managers/room.py | 78 +++++++++++++------ .../server/handlers/managers/service.py | 2 +- .../server/handlers/websocket_server.py | 24 ++++-- dragonion_server/modules/server/server.py | 12 +++ pyproject.toml | 2 +- 9 files changed, 125 insertions(+), 51 deletions(-) create mode 100644 dragonion_server/modules/server/handlers/managers/exceptions.py diff --git a/dragonion_server/cli/cmd/service/run.py b/dragonion_server/cli/cmd/service/run.py index e44f238..3ff57c4 100644 --- a/dragonion_server/cli/cmd/service/run.py +++ b/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) diff --git a/dragonion_server/modules/server/__init__.py b/dragonion_server/modules/server/__init__.py index 7e031d0..9ac0db8 100644 --- a/dragonion_server/modules/server/__init__.py +++ b/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' ] diff --git a/dragonion_server/modules/server/handlers/managers/connection.py b/dragonion_server/modules/server/handlers/managers/connection.py index 8f65b1d..69490e3 100644 --- a/dragonion_server/modules/server/handlers/managers/connection.py +++ b/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, - WebUserMessage + WebErrorMessage ) @@ -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 - ) - ) diff --git a/dragonion_server/modules/server/handlers/managers/exceptions.py b/dragonion_server/modules/server/handlers/managers/exceptions.py new file mode 100644 index 0000000..98f9fab --- /dev/null +++ b/dragonion_server/modules/server/handlers/managers/exceptions.py @@ -0,0 +1,2 @@ +class GotInvalidWebmessage(Exception): + pass diff --git a/dragonion_server/modules/server/handlers/managers/room.py b/dragonion_server/modules/server/handlers/managers/room.py index 8e3a7b0..622ea22 100644 --- a/dragonion_server/modules/server/handlers/managers/room.py +++ b/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 + WebErrorMessage, + 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 - await connection.send_connect() - print(f'Accepted {username}') + self.connections[connection_message.username] = connection + await connection.send_webmessage(WebConnectionResultMessage( + 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 message: content + :param broadcastable: String object with json representation of + WebBroadcastableMessage :return: """ - await self.broadcast_webmessage( - WebMessageMessage( - username=from_username, - message=message - ) - ) + try: + for to_username in broadcastable.messages.keys(): + try: + 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): """ @@ -96,11 +131,10 @@ class Room(object): """ Broadcasts that user is disconnected :param username: Username of user that disconnected - :return: + :return: """ await self.broadcast_webmessage( - WebUserMessage( - type="disconnect", + WebDisconnectMessage( username=username ) ) diff --git a/dragonion_server/modules/server/handlers/managers/service.py b/dragonion_server/modules/server/handlers/managers/service.py index 3038902..d375793 100644 --- a/dragonion_server/modules/server/handlers/managers/service.py +++ b/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 ) diff --git a/dragonion_server/modules/server/handlers/websocket_server.py b/dragonion_server/modules/server/handlers/websocket_server.py index f8ad9c8..c472bf7 100644 --- a/dragonion_server/modules/server/handlers/websocket_server.py +++ b/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.from_json(data) + 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) + 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: diff --git a/dragonion_server/modules/server/server.py b/dragonion_server/modules/server/server.py index 2f9cdd1..ac719b9 100644 --- a/dragonion_server/modules/server/server.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 120663f..a96f899 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ authors = ["BarsTiger "] 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"