119 lines
3.3 KiB
Python
119 lines
3.3 KiB
Python
from fastapi import (
|
|
FastAPI,
|
|
Request,
|
|
Response,
|
|
WebSocket,
|
|
WebSocketException,
|
|
status,
|
|
Cookie,
|
|
Depends,
|
|
)
|
|
from fastapi.staticfiles import StaticFiles
|
|
from kennel.engine.systems.system import SystemType
|
|
from kennel.engine.systems.network import Event, Publishable, EventType
|
|
from typing import Annotated, Optional
|
|
from .kennel import (
|
|
kennel,
|
|
system_manager,
|
|
entity_manager,
|
|
create_session_controllable_entities,
|
|
)
|
|
from .app import app, templates, logger
|
|
from .kennelcats import KennelCatService
|
|
from .middleware import logger_middleware
|
|
from .config import config
|
|
import asyncio
|
|
import uuid
|
|
|
|
app.mount("/static", StaticFiles(directory="static"), name="static")
|
|
|
|
loop = asyncio.get_event_loop()
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
logger.info("Starting Kennel...")
|
|
loop.create_task(kennel.run())
|
|
|
|
|
|
@app.on_event("shutdown")
|
|
async def shutdown_event():
|
|
logger.info("Stopping Kennel...")
|
|
kennel.stop()
|
|
loop.stop()
|
|
logger.info("Kennel stopped")
|
|
|
|
|
|
@app.get("/")
|
|
async def index(request: Request):
|
|
return templates.TemplateResponse(request=request, name="index.html")
|
|
|
|
|
|
@app.get("/assign")
|
|
async def assign(
|
|
response: Response,
|
|
session: Annotated[Optional[str], Cookie()] = None,
|
|
):
|
|
if session is None:
|
|
session = str(uuid.uuid4().hex)
|
|
response.set_cookie(key="session", value=session, max_age=config.COOKIE_MAX_AGE)
|
|
return {"session": session}
|
|
|
|
|
|
async def get_cookie_or_token(
|
|
websocket: WebSocket,
|
|
session: Annotated[str | None, Cookie()] = None,
|
|
):
|
|
if session is None:
|
|
raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
|
|
return session
|
|
|
|
|
|
class WebSocketClient(Publishable):
|
|
def __init__(self, websocket: WebSocket):
|
|
self.websocket = websocket
|
|
|
|
async def publish(self, event: Event):
|
|
await self.websocket.send_json(event.dict())
|
|
|
|
|
|
@app.websocket("/ws")
|
|
async def websocket_endpoint(
|
|
websocket: WebSocket,
|
|
session: Annotated[str, Depends(get_cookie_or_token)],
|
|
):
|
|
await websocket.accept()
|
|
logger.info(f"Websocket connection established for session {session}")
|
|
|
|
session_entities = create_session_controllable_entities(session)
|
|
try:
|
|
network_system = system_manager.get_system(SystemType.NETWORK)
|
|
if network_system is None:
|
|
raise "Network system not found"
|
|
|
|
for entity in session_entities:
|
|
logger.info(f"Adding entity {entity.id} for session {session}")
|
|
entity_manager.add_entity(entity)
|
|
network_system.add_event(Event(EventType.ENTITY_BORN, {"id": entity.id}))
|
|
|
|
network_system.add_client(session, WebSocketClient(websocket))
|
|
while True:
|
|
message = await websocket.receive_json()
|
|
network_system.client_event(session, Event.from_dict(message))
|
|
except Exception as e:
|
|
logger.error(f"WebSocket exception {e}")
|
|
finally:
|
|
logger.info("Websocket connection closed")
|
|
for entity in session_entities:
|
|
logger.info(f"Removing entity {entity.id} for session {session}")
|
|
entity_manager.remove_entity(entity.id)
|
|
network_system.add_event(Event(EventType.ENTITY_DEATH, {"id": entity.id}))
|
|
|
|
network_system.remove_client(session)
|
|
await websocket.close()
|
|
|
|
|
|
@app.get("/healthcheck")
|
|
async def healthcheck():
|
|
return Response("healthy")
|