2024-02-29 23:42:12 -05:00
# pylint: disable=cyclic-import
2024-02-29 23:26:24 -05:00
import json
import re
2024-08-26 17:46:42 -04:00
from pathlib import Path
from typing import Optional , Tuple , Union
2024-02-29 23:26:24 -05:00
import aiohttp
import discord
import websockets
2024-03-01 14:52:14 -05:00
from pydactyl import PterodactylClient
2024-08-26 17:46:42 -04:00
from redbot . core . data_manager import bundled_data_path
2024-03-02 15:38:06 -05:00
from redbot . core . utils . chat_formatting import bold , pagify
2024-02-29 23:26:24 -05:00
from pterodactyl . config import config
2024-03-07 01:32:27 -05:00
from pterodactyl . logger import logger , websocket_logger
2024-02-29 23:26:24 -05:00
from pterodactyl . pterodactyl import Pterodactyl
async def establish_websocket_connection ( coginstance : Pterodactyl ) - > None :
2024-04-08 06:24:53 -04:00
await coginstance . bot . wait_until_red_ready ( )
2024-02-29 23:42:12 -05:00
base_url = await config . base_url ( )
2024-03-01 15:44:26 -05:00
base_url = base_url [ : - 1 ] if base_url . endswith ( ' / ' ) else base_url
2024-02-29 23:42:12 -05:00
2024-03-01 13:16:22 -05:00
logger . info ( " Establishing WebSocket connection " )
2024-02-29 23:42:12 -05:00
websocket_credentials = await retrieve_websocket_credentials ( coginstance )
2024-03-07 01:32:27 -05:00
async with websockets . connect ( websocket_credentials [ ' data ' ] [ ' socket ' ] , origin = base_url , ping_timeout = 60 , logger = websocket_logger ) as websocket :
2024-02-29 23:42:12 -05:00
logger . info ( " WebSocket connection established " )
auth_message = json . dumps ( { " event " : " auth " , " args " : [ websocket_credentials [ ' data ' ] [ ' token ' ] ] } )
await websocket . send ( auth_message )
2024-03-01 13:16:22 -05:00
logger . info ( " Authentication message sent " )
2024-02-29 23:42:12 -05:00
coginstance . websocket = websocket
2024-02-29 23:52:48 -05:00
while True : # pylint: disable=too-many-nested-blocks
2024-03-03 02:01:01 -05:00
message = json . loads ( await websocket . recv ( ) )
if message [ ' event ' ] in ( ' token expiring ' , ' token expired ' ) :
2024-03-01 13:16:22 -05:00
logger . info ( " Received token expiring/expired event. Refreshing token. " )
2024-02-29 23:52:48 -05:00
websocket_credentials = await retrieve_websocket_credentials ( coginstance )
auth_message = json . dumps ( { " event " : " auth " , " args " : [ websocket_credentials [ ' data ' ] [ ' token ' ] ] } )
await websocket . send ( auth_message )
2024-03-01 13:16:22 -05:00
logger . info ( " Authentication message sent " )
2024-02-29 23:52:48 -05:00
2024-03-03 02:01:01 -05:00
if message [ ' event ' ] == ' auth success ' :
2024-02-29 23:52:48 -05:00
logger . info ( " WebSocket authentication successful " )
2024-03-03 02:01:01 -05:00
if message [ ' event ' ] == ' console output ' and await config . console_channel ( ) is not None :
2024-03-02 19:43:28 -05:00
regex_blacklist : dict = await config . regex_blacklist ( )
2024-03-03 02:01:01 -05:00
matches = [ re . search ( regex , message [ ' args ' ] [ 0 ] ) for regex in regex_blacklist . values ( ) ]
2024-03-02 19:50:48 -05:00
2024-03-02 20:10:16 -05:00
if await config . current_status ( ) in ( ' running ' , ' ' ) and not any ( matches ) :
2024-03-03 02:01:01 -05:00
content = remove_ansi_escape_codes ( message [ ' args ' ] [ 0 ] )
2024-03-01 22:38:49 -05:00
if await config . mask_ip ( ) is True :
content = mask_ip ( content )
2024-02-29 23:52:48 -05:00
2024-03-09 23:26:27 -05:00
console_channel = coginstance . bot . get_channel ( await config . console_channel ( ) )
chat_channel = coginstance . bot . get_channel ( await config . chat_channel ( ) )
if console_channel is not None :
2024-02-29 23:52:48 -05:00
if content . startswith ( ' [ ' ) :
pagified_content = pagify ( content , delims = [ " " , " \n " ] )
for page in pagified_content :
2024-03-09 23:26:27 -05:00
await console_channel . send ( content = page , allowed_mentions = discord . AllowedMentions . none ( ) )
2024-02-29 23:52:48 -05:00
server_message = await check_if_server_message ( content )
if server_message :
2024-03-09 23:26:27 -05:00
if chat_channel is not None :
await chat_channel . send ( server_message if len ( server_message ) < 2000 else server_message [ : 1997 ] + ' ... ' , allowed_mentions = discord . AllowedMentions . none ( ) )
2024-02-29 23:52:48 -05:00
2024-03-02 15:27:34 -05:00
chat_message = await check_if_chat_message ( content )
if chat_message :
info = await get_info ( chat_message [ ' username ' ] )
if info is not None :
await send_chat_discord ( coginstance , chat_message [ ' username ' ] , chat_message [ ' message ' ] , info [ ' data ' ] [ ' player ' ] [ ' avatar ' ] )
else :
await send_chat_discord ( coginstance , chat_message [ ' username ' ] , chat_message [ ' message ' ] , ' https://seafsh.cc/u/j3AzqQ.png ' )
2024-03-02 15:20:33 -05:00
2024-03-01 00:23:00 -05:00
join_message = await check_if_join_message ( content )
if join_message :
2024-03-09 23:26:27 -05:00
if chat_channel is not None :
if coginstance . bot . embed_requested ( chat_channel ) :
2024-08-26 19:43:57 -04:00
embed , img = await generate_join_leave_embed ( coginstance = coginstance , username = join_message , join = True )
2024-08-26 17:46:42 -04:00
if img :
with open ( img , ' rb ' ) as file :
await chat_channel . send ( embed = embed , file = file )
else :
await chat_channel . send ( embed = embed )
2024-03-01 00:23:00 -05:00
else :
2024-03-09 23:26:27 -05:00
await chat_channel . send ( f " { join_message } joined the game " , allowed_mentions = discord . AllowedMentions . none ( ) )
2024-03-01 00:23:00 -05:00
leave_message = await check_if_leave_message ( content )
if leave_message :
2024-03-09 23:26:27 -05:00
if chat_channel is not None :
if coginstance . bot . embed_requested ( chat_channel ) :
2024-08-26 19:43:57 -04:00
embed , img = await generate_join_leave_embed ( coginstance = coginstance , username = leave_message , join = False )
2024-08-26 17:46:42 -04:00
if img :
with open ( img , ' rb ' ) as file :
await chat_channel . send ( embed = embed , file = file )
else :
await chat_channel . send ( embed = embed )
2024-03-01 00:23:00 -05:00
else :
2024-03-09 23:26:27 -05:00
await chat_channel . send ( f " { leave_message } left the game " , allowed_mentions = discord . AllowedMentions . none ( ) )
2024-03-01 00:23:00 -05:00
2024-03-01 00:46:51 -05:00
achievement_message = await check_if_achievement_message ( content )
if achievement_message :
2024-03-09 23:26:27 -05:00
if chat_channel is not None :
if coginstance . bot . embed_requested ( chat_channel ) :
await chat_channel . send ( embed = await generate_achievement_embed ( achievement_message [ ' username ' ] , achievement_message [ ' achievement ' ] , achievement_message [ ' challenge ' ] ) )
2024-03-01 00:46:51 -05:00
else :
2024-03-09 23:26:27 -05:00
await chat_channel . send ( f " { achievement_message [ ' username ' ] } has { ' completed the challenge ' if achievement_message [ ' challenge ' ] else ' made the advancement ' } { achievement_message [ ' achievement ' ] } " )
2024-03-01 00:46:51 -05:00
2024-03-03 02:01:33 -05:00
if message [ ' event ' ] == ' status ' :
2024-03-01 00:02:42 -05:00
old_status = await config . current_status ( )
2024-03-03 02:01:33 -05:00
current_status = message [ ' args ' ] [ 0 ]
2024-03-01 00:02:42 -05:00
if old_status != current_status :
await config . current_status . set ( current_status )
if await config . console_channel ( ) is not None :
console = coginstance . bot . get_channel ( await config . console_channel ( ) )
if console is not None :
await console . send ( f " Server status changed! ` { current_status } ` " )
if await config . chat_channel ( ) is not None :
if current_status == ' running ' and await config . startup_msg ( ) is not None :
chat = coginstance . bot . get_channel ( await config . chat_channel ( ) )
if chat is not None :
await chat . send ( await config . startup_msg ( ) )
if current_status == ' stopping ' and await config . shutdown_msg ( ) is not None :
chat = coginstance . bot . get_channel ( await config . chat_channel ( ) )
if chat is not None :
await chat . send ( await config . shutdown_msg ( ) )
2024-02-29 23:42:12 -05:00
async def retrieve_websocket_credentials ( coginstance : Pterodactyl ) - > Optional [ dict ] :
2024-03-01 01:17:31 -05:00
pterodactyl_keys = await coginstance . bot . get_shared_api_tokens ( " pterodactyl " )
api_key = pterodactyl_keys . get ( " api_key " )
if api_key is None :
coginstance . task . cancel ( )
2024-03-01 14:52:14 -05:00
raise ValueError ( " Pterodactyl API key not set. Please set it using `[p]set api`. " )
2024-02-29 23:26:24 -05:00
base_url = await config . base_url ( )
2024-03-01 01:17:31 -05:00
if base_url is None :
coginstance . task . cancel ( )
2024-03-01 14:52:14 -05:00
raise ValueError ( " Pterodactyl base URL not set. Please set it using `[p]pterodactyl config url`. " )
2024-02-29 23:26:24 -05:00
server_id = await config . server_id ( )
2024-03-01 01:17:31 -05:00
if server_id is None :
coginstance . task . cancel ( )
2024-03-01 14:52:14 -05:00
raise ValueError ( " Pterodactyl server ID not set. Please set it using `[p]pterodactyl config serverid`. " )
2024-03-07 02:39:49 -05:00
client = PterodactylClient ( base_url , api_key ) . client
2024-03-01 14:52:14 -05:00
coginstance . client = client
websocket_credentials = client . servers . get_websocket ( server_id )
logger . debug ( """ Websocket connection details retrieved:
Socket : % s
Token : % s . . . """ ,
websocket_credentials [ ' data ' ] [ ' socket ' ] ,
websocket_credentials [ ' data ' ] [ ' token ' ] [ : 20 ]
)
return websocket_credentials
#NOTE - The token is truncated to prevent it from being logged in its entirety, for security reasons
2024-02-29 23:26:24 -05:00
def remove_ansi_escape_codes ( text : str ) - > str :
ansi_escape = re . compile ( r ' \ x1B(?:[@-Z \\ -_]| \ [[0-?]*[ -/]*[@-~]) ' )
#NOTE - https://chat.openai.com/share/d92f9acf-d776-4fd6-a53f-b14ac15dd540
return ansi_escape . sub ( ' ' , text )
2024-02-29 23:50:54 -05:00
async def check_if_server_message ( text : str ) - > Union [ bool , str ] :
regex = await config . server_regex ( )
match : Optional [ re . Match [ str ] ] = re . match ( regex , text )
if match :
2024-08-26 19:43:57 -04:00
logger . trace ( " Message is a server message " )
2024-03-02 12:41:54 -05:00
return match . group ( 1 )
2024-02-29 23:50:54 -05:00
return False
2024-02-29 23:26:24 -05:00
async def check_if_chat_message ( text : str ) - > Union [ bool , dict ] :
regex = await config . chat_regex ( )
match : Optional [ re . Match [ str ] ] = re . match ( regex , text )
if match :
2024-03-02 15:11:40 -05:00
groups = { " username " : match . group ( 1 ) , " message " : match . group ( 2 ) }
2024-08-26 19:43:57 -04:00
logger . trace ( " Message is a chat message \n %s " , json . dumps ( groups ) )
2024-02-29 23:26:24 -05:00
return groups
return False
2024-03-01 00:23:00 -05:00
async def check_if_join_message ( text : str ) - > Union [ bool , str ] :
regex = await config . join_regex ( )
match : Optional [ re . Match [ str ] ] = re . match ( regex , text )
if match :
2024-08-26 19:43:57 -04:00
logger . trace ( " Message is a join message " )
2024-03-01 00:23:00 -05:00
return match . group ( 1 )
return False
async def check_if_leave_message ( text : str ) - > Union [ bool , str ] :
regex = await config . leave_regex ( )
match : Optional [ re . Match [ str ] ] = re . match ( regex , text )
if match :
2024-08-26 19:43:57 -04:00
logger . trace ( " Message is a leave message " )
2024-03-01 00:23:00 -05:00
return match . group ( 1 )
return False
2024-03-01 00:46:51 -05:00
async def check_if_achievement_message ( text : str ) - > Union [ bool , dict ] :
regex = await config . achievement_regex ( )
match : Optional [ re . Match [ str ] ] = re . match ( regex , text )
if match :
groups = { " username " : match . group ( 1 ) , " achievement " : match . group ( 3 ) }
if match . group ( 2 ) == " completed the challenge " :
groups [ " challenge " ] = True
else :
groups [ " challenge " ] = False
2024-08-26 19:43:57 -04:00
logger . trace ( " Message is an achievement message " )
2024-03-01 00:46:51 -05:00
return groups
return False
2024-02-29 23:26:24 -05:00
async def get_info ( username : str ) - > Optional [ dict ] :
2024-08-26 19:43:57 -04:00
logger . verbose ( " Retrieving player info for %s " , username )
2024-02-29 23:26:24 -05:00
endpoint = await config . api_endpoint ( )
async with aiohttp . ClientSession ( ) as session :
async with session . get ( f " https://playerdb.co/api/player/ { endpoint } / { username } " ) as response :
if response . status == 200 :
2024-08-26 19:43:57 -04:00
logger . verbose ( " Player info retrieved for %s " , username )
2024-02-29 23:26:24 -05:00
return await response . json ( )
2024-08-26 19:43:57 -04:00
logger . warning ( " Failed to retrieve player info for %s : %s " , username , response . status )
2024-02-29 23:26:24 -05:00
return None
async def send_chat_discord ( coginstance : Pterodactyl , username : str , message : str , avatar_url : str ) - > None :
2024-08-26 19:43:57 -04:00
logger . trace ( " Sending chat message to Discord " )
2024-02-29 23:26:24 -05:00
channel = coginstance . bot . get_channel ( await config . chat_channel ( ) )
if channel is not None :
webhooks = await channel . webhooks ( )
webhook = discord . utils . get ( webhooks , name = " Pterodactyl Chat " )
if webhook is None :
webhook = await channel . create_webhook ( name = " Pterodactyl Chat " )
2024-03-04 19:53:12 -05:00
await webhook . send ( content = message , username = username , avatar_url = avatar_url , allowed_mentions = discord . AllowedMentions ( everyone = False , roles = False , users = True ) )
2024-08-26 19:43:57 -04:00
logger . trace ( " Chat message sent to Discord " )
2024-02-29 23:26:24 -05:00
else :
2024-03-07 02:19:00 -05:00
logger . warning ( " Chat channel not set. Skipping sending chat message to Discord " )
2024-03-01 00:23:00 -05:00
2024-08-26 17:46:42 -04:00
async def generate_join_leave_embed ( coginstance : Pterodactyl , username : str , join : bool ) - > Tuple [ discord . Embed , Optional [ Union [ str , Path ] ] ] :
2024-03-01 00:23:00 -05:00
embed = discord . Embed ( )
2024-03-01 00:25:39 -05:00
embed . color = discord . Color . green ( ) if join else discord . Color . red ( )
2024-03-01 00:23:00 -05:00
embed . description = await config . join_msg ( ) if join else await config . leave_msg ( )
info = await get_info ( username )
if info :
2024-08-26 17:46:42 -04:00
img = None
2024-03-01 00:23:00 -05:00
embed . set_author ( name = username , icon_url = info [ ' data ' ] [ ' player ' ] [ ' avatar ' ] )
else :
2024-08-26 17:46:42 -04:00
img = bundled_data_path ( coginstance ) / " unknown.png "
embed . set_author ( name = username , icon_url = ' attachment://unknown.png ' )
2024-03-01 00:23:00 -05:00
embed . timestamp = discord . utils . utcnow ( )
2024-08-26 17:46:42 -04:00
return embed , img
2024-03-01 00:46:51 -05:00
2024-08-26 17:46:42 -04:00
async def generate_achievement_embed ( coginstance : Pterodactyl , username : str , achievement : str , challenge : bool ) - > Tuple [ discord . Embed , Optional [ Union [ str , Path ] ] ] :
2024-03-01 00:46:51 -05:00
embed = discord . Embed ( )
2024-03-02 15:42:29 -05:00
embed . color = discord . Color . from_str ( ' #a800a7 ' ) if challenge else discord . Color . from_str ( ' #54fb54 ' )
2024-03-02 15:38:06 -05:00
embed . description = f " { bold ( username ) } has { ' completed the challenge ' if challenge else ' made the advancement ' } { bold ( achievement ) } "
2024-03-01 00:46:51 -05:00
info = await get_info ( username )
if info :
2024-08-26 17:46:42 -04:00
img = None
2024-03-01 00:46:51 -05:00
embed . set_author ( name = username , icon_url = info [ ' data ' ] [ ' player ' ] [ ' avatar ' ] )
else :
2024-08-26 17:46:42 -04:00
img = bundled_data_path ( coginstance ) / " unknown.png "
embed . set_author ( name = username , icon_url = ' attachment://unknown.png ' )
2024-03-01 00:46:51 -05:00
embed . timestamp = discord . utils . utcnow ( )
2024-08-26 17:46:42 -04:00
return embed , img
2024-03-01 22:38:49 -05:00
def mask_ip ( string : str ) - > str :
2024-03-04 23:12:32 -05:00
def check ( match : re . Match [ str ] ) :
2024-03-01 22:38:49 -05:00
ip = match . group ( 0 )
2024-03-01 22:39:58 -05:00
masked_ip = ' . ' . join ( r ' \ * ' * len ( octet ) for octet in ip . split ( ' . ' ) )
2024-03-01 22:38:49 -05:00
return masked_ip
2024-03-01 23:43:46 -05:00
return re . sub ( r ' \ b(?: \ d { 1,3} \ .) {3} \ d { 1,3} \ b ' , check , string )