152 строки
3.6 KiB
Python
152 строки
3.6 KiB
Python
import asyncio
|
|
import hashlib
|
|
import os
|
|
|
|
import aiohttp
|
|
from fastapi import FastAPI, Query, WebSocket, WebSocketDisconnect
|
|
from fastapi.exceptions import HTTPException
|
|
from fastapi.responses import Response, HTMLResponse, FileResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
from jinja2 import Environment, FileSystemLoader
|
|
|
|
import config
|
|
|
|
|
|
app = FastAPI()
|
|
app.mount(
|
|
path='/static',
|
|
app=StaticFiles(
|
|
directory=os.path.join(
|
|
config.Main.working_directory,
|
|
'static',
|
|
),
|
|
),
|
|
)
|
|
|
|
env = Environment(
|
|
loader=FileSystemLoader(
|
|
searchpath=os.path.join(
|
|
'.',
|
|
'templates',
|
|
),
|
|
),
|
|
enable_async=True,
|
|
)
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.connections: list[WebSocket] = list[WebSocket]()
|
|
|
|
async def connect(self, websocket: WebSocket):
|
|
await websocket.accept()
|
|
self.connections.append(websocket)
|
|
|
|
def disconnect(self, websocket: WebSocket):
|
|
self.connections.remove(websocket)
|
|
|
|
async def broadcast(self, data: dict):
|
|
for connection in self.connections:
|
|
asyncio.ensure_future(connection.send_json(data))
|
|
|
|
|
|
overlay_manager = ConnectionManager()
|
|
plugin_manager = ConnectionManager()
|
|
|
|
|
|
@app.options(
|
|
path='/api/v1/overlay',
|
|
)
|
|
async def function():
|
|
return Response(
|
|
headers={
|
|
'Access-Control-Allow-Origin': '*',
|
|
'Access-Control-Allow-Methods': '*',
|
|
'Access-Control-Allow-Headers': '*',
|
|
},
|
|
status_code=200,
|
|
)
|
|
|
|
|
|
@app.get(
|
|
path='/overlay',
|
|
)
|
|
async def function(
|
|
justify_content: str = Query(
|
|
default='end',
|
|
alias='justify-content',
|
|
),
|
|
align_items: str = Query(
|
|
default='end',
|
|
alias='align-items',
|
|
),
|
|
):
|
|
template = env.get_template('overlay.jinja2')
|
|
return HTMLResponse(
|
|
content=await template.render_async(
|
|
justify_content=justify_content,
|
|
align_items=align_items,
|
|
),
|
|
status_code=200,
|
|
)
|
|
|
|
|
|
@app.get(
|
|
path='/image',
|
|
)
|
|
async def function(
|
|
src: str = Query(
|
|
default=...,
|
|
),
|
|
):
|
|
image_name = hashlib.sha1(src.encode('ascii')).hexdigest()
|
|
image_path = os.path.join(
|
|
config.Main.images_directory,
|
|
image_name,
|
|
)
|
|
if not os.path.exists(image_path):
|
|
os.makedirs(
|
|
name=config.Main.images_directory,
|
|
exist_ok=True,
|
|
)
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(src) as response:
|
|
if not response.ok:
|
|
raise HTTPException(
|
|
status_code=response.status,
|
|
)
|
|
with open(
|
|
file=image_path,
|
|
mode='wb',
|
|
) as file:
|
|
file.write(await response.read())
|
|
return FileResponse(
|
|
path=image_path,
|
|
)
|
|
|
|
|
|
@app.websocket('/ws/v1/overlay')
|
|
async def function(
|
|
websocket: WebSocket,
|
|
):
|
|
await overlay_manager.connect(websocket)
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
await plugin_manager.broadcast(data)
|
|
except WebSocketDisconnect:
|
|
overlay_manager.disconnect(websocket)
|
|
|
|
|
|
@app.websocket('/ws/v1/plugin')
|
|
async def function(
|
|
websocket: WebSocket,
|
|
):
|
|
await plugin_manager.connect(websocket)
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
await overlay_manager.broadcast(data)
|
|
except WebSocketDisconnect:
|
|
plugin_manager.disconnect(websocket)
|