98 lines
2.1 KiB
Python
98 lines
2.1 KiB
Python
|
import enum
|
||
|
|
import json
|
||
|
|
import asyncio
|
||
|
|
import websockets
|
||
|
|
from typing import Union, Any, Dict, Callable, Coroutine
|
||
|
|
|
||
|
|
from .atconnect import REQUA
|
||
|
|
|
||
|
|
class Streams(enum.IntEnum):
|
||
|
|
status = 0
|
||
|
|
queue = 1
|
||
|
|
console = 2
|
||
|
|
ram = 3
|
||
|
|
tps = 4
|
||
|
|
|
||
|
|
class AternosWss:
|
||
|
|
|
||
|
|
def __init__(self, session:str, servid:str) -> None:
|
||
|
|
|
||
|
|
self.session = session
|
||
|
|
self.servid = servid
|
||
|
|
self.recv = {}
|
||
|
|
|
||
|
|
def wssreceiver(self, stream:int) -> Callable[[Callable[[Any],Coroutine[Any,Any,None]]],Any]:
|
||
|
|
def decorator(func:Callable[[Any],Coroutine[Any,Any,None]]) -> None:
|
||
|
|
self.recv[stream] = func
|
||
|
|
return decorator
|
||
|
|
|
||
|
|
async def connect(self) -> None:
|
||
|
|
|
||
|
|
headers = [
|
||
|
|
('User-Agent', REQUA),
|
||
|
|
(
|
||
|
|
'Cookie',
|
||
|
|
f'ATERNOS_SESSION={self.session}; ' +\
|
||
|
|
f'ATERNOS_SERVER={self.servid}'
|
||
|
|
)
|
||
|
|
]
|
||
|
|
self.socket = await websockets.connect(
|
||
|
|
'wss://aternos.org/hermes',
|
||
|
|
extra_headers=headers
|
||
|
|
)
|
||
|
|
asyncio.run(wssworker())
|
||
|
|
|
||
|
|
async def close(self) -> None:
|
||
|
|
|
||
|
|
await self.socket.close()
|
||
|
|
del self.socket
|
||
|
|
|
||
|
|
async def send(self, obj:Union[Dict[str, Any],str]) -> None:
|
||
|
|
|
||
|
|
if isinstance(obj, dict):
|
||
|
|
obj = json.dumps(obj)
|
||
|
|
|
||
|
|
self.socket.send(obj)
|
||
|
|
|
||
|
|
async def wssworker(self) -> None:
|
||
|
|
|
||
|
|
keep = asyncio.create_task(keepalive())
|
||
|
|
msgs = asyncio.create_task(receiver())
|
||
|
|
await keep
|
||
|
|
await msgs
|
||
|
|
|
||
|
|
async def keepalive(self) -> None:
|
||
|
|
|
||
|
|
while True:
|
||
|
|
await asyncio.sleep(49)
|
||
|
|
await self.socket.send('{"type":"\u2764"}')
|
||
|
|
|
||
|
|
async def receiver(self) -> None:
|
||
|
|
|
||
|
|
while True:
|
||
|
|
data = await self.socket.recv()
|
||
|
|
obj = json.loads(data)
|
||
|
|
|
||
|
|
if obj['type'] == 'line':
|
||
|
|
msgtype = Streams.console
|
||
|
|
msg = obj['data'].strip('\r\n ')
|
||
|
|
|
||
|
|
elif obj['type'] == 'heap':
|
||
|
|
msgtype = Streams.ram
|
||
|
|
msg = int(obj['data']['usage'])
|
||
|
|
|
||
|
|
elif obj['type'] == 'tick':
|
||
|
|
msgtype = Streams.tps
|
||
|
|
ticks = 1000 / obj['data']['averageTickTime']
|
||
|
|
msg = 20 if ticks > 20 else ticks
|
||
|
|
|
||
|
|
elif obj['type'] == 'status':
|
||
|
|
msgtype = Streams.status
|
||
|
|
msg = json.loads(obj['message'])
|
||
|
|
|
||
|
|
if msgtype in self.recv:
|
||
|
|
t = asyncio.create_task(
|
||
|
|
self.recv[msgtype](msg)
|
||
|
|
)
|
||
|
|
await t
|