hack-house/cmd_chat/server/server.py

88 lines
2.4 KiB
Python

import asyncio
import rsa
from cryptography.fernet import Fernet
from functools import partial
from sanic.worker.loader import AppLoader
from sanic.response import HTTPResponse
from sanic import Sanic, Request, response, Websocket
from cmd_chat.server.models import Message
from cmd_chat.server.services import (
_get_bytes_and_serialize,
_check_ws_for_close_status,
_generate_new_message,
_generate_update_payload
)
app = Sanic("app")
app.config.OAS = False
# Message structure is:
# [username: message, ...]
MESSAGES_MEMORY_DB: list[Message] = []
# Users structure is
# {Ip, Username: Public key}
USERS: dict[str, str] = {}
PUBLIC_KEY = Fernet.generate_key()
def attach_endpoints(app: Sanic):
@app.websocket("/talk")
async def talk_ws_view(request: Request, ws: Websocket) -> HTTPResponse:
while True:
serialized_message: dict = await _get_bytes_and_serialize(ws)
await _check_ws_for_close_status(
serialized_message,
ws
)
new_message = await _generate_new_message(
serialized_message.get("text")
)
MESSAGES_MEMORY_DB.append(new_message)
await ws.send(
str({"status": "ok"})
)
await asyncio.sleep(0.2)
@app.websocket("/update")
async def update_ws_view(request: Request, ws: Websocket) -> HTTPResponse:
while True:
payload = await _generate_update_payload(
MESSAGES_MEMORY_DB,
USERS
)
await ws.send(payload.encode())
await asyncio.sleep(0.2)
@app.route('/get_key', methods=['GET', 'POST'])
async def get_key_view(request: Request) -> HTTPResponse:
public_key = rsa.PublicKey.load_pkcs1(request.form.get('pubkey'))
encrypted_data = rsa.encrypt(PUBLIC_KEY, public_key)
if request.ip not in USERS:
USERS[f"{request.ip}, {request.form.get('username')}"] = PUBLIC_KEY
return response.raw(encrypted_data)
def create_app(app_name: str) -> Sanic:
app = Sanic(app_name)
attach_endpoints(app)
return app
def run_server(host: str, port: int, dev: bool=False) -> None:
loader = AppLoader(factory=partial(create_app, "CMD_SERVER"))
app = loader.load()
app.prepare(host=host, port=port, dev=dev)
Sanic.serve(primary=app, app_loader=loader)