fix(pterodactyl): handle disconnecting from the websocket

This commit is contained in:
Seaswimmer 2024-02-28 11:46:00 -05:00
parent e6ca4071f3
commit d391aeeb7b
Signed by untrusted user: cswimr
GPG key ID: B8953EC01E5C4063

View file

@ -47,7 +47,8 @@ class Pterodactyl(commands.Cog):
self.logger.error('Failed to retrieve Pterodactyl websocket: %s', e) self.logger.error('Failed to retrieve Pterodactyl websocket: %s', e)
return return
async with websockets.connect(websocket_credentials['data']['socket'], extra_headers=extra_headers) as websocket: async for websocket in websockets.connect(websocket_credentials['data']['socket'], extra_headers=extra_headers):
try:
self.logger.debug("WebSocket connection established") self.logger.debug("WebSocket connection established")
auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]}) auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]})
@ -68,6 +69,10 @@ class Pterodactyl(commands.Cog):
if json.loads(message)['event'] == 'auth success': if json.loads(message)['event'] == 'auth success':
self.logger.debug("Authentication successful") self.logger.debug("Authentication successful")
self.logger.debug("Received message: %s", message) self.logger.debug("Received message: %s", message)
except websockets.exceptions.ConnectionClosed as e:
self.logger.debug("WebSocket connection closed: %s", e)
websocket_credentials = client.servers.get_websocket(server_id)
continue
def get_task(self): def get_task(self):
return self.bot.loop.create_task(self.establish_websocket_connection(), name="Pterodactyl Websocket Connection") return self.bot.loop.create_task(self.establish_websocket_connection(), name="Pterodactyl Websocket Connection")