Working development server (no encryption now)
This commit is contained in:
BarsTiger
1
dragonion_server/modules/server/handlers/__init__.py
Normal file
1
dragonion_server/modules/server/handlers/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
pass
|
||||
@@ -0,0 +1 @@
|
||||
pass
|
||||
@@ -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
|
||||
)
|
||||
)
|
||||
95
dragonion_server/modules/server/handlers/managers/room.py
Normal file
95
dragonion_server/modules/server/handlers/managers/room.py
Normal file
@@ -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
|
||||
52
dragonion_server/modules/server/handlers/managers/service.py
Normal file
52
dragonion_server/modules/server/handlers/managers/service.py
Normal file
@@ -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}'
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
pass
|
||||
@@ -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)
|
||||
44
dragonion_server/modules/server/handlers/websocket_server.py
Normal file
44
dragonion_server/modules/server/handlers/websocket_server.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user