From 306148ea6927c7d228470b4977a12e5cb6ab901a Mon Sep 17 00:00:00 2001
From: SeaswimmerTheFsh <seaswimmerthefsh@gmail.com>
Date: Fri, 1 Mar 2024 22:38:49 -0500
Subject: [PATCH] feat(pterodactyl): added masking of IP addresses in the
 console output

---
 pterodactyl/config.py      | 1 +
 pterodactyl/pterodactyl.py | 9 +++++++++
 pterodactyl/websocket.py   | 9 +++++++++
 3 files changed, 19 insertions(+)

diff --git a/pterodactyl/config.py b/pterodactyl/config.py
index a116375..fec0336 100644
--- a/pterodactyl/config.py
+++ b/pterodactyl/config.py
@@ -22,4 +22,5 @@ def register_config(config_obj: Config) -> None:
         shutdown_msg='Server stopped!',
         join_msg='Welcome to the server! 👋',
         leave_msg='Goodbye! 👋',
+        mask_ip=True,
     )
diff --git a/pterodactyl/pterodactyl.py b/pterodactyl/pterodactyl.py
index b099c39..0c6ad27 100644
--- a/pterodactyl/pterodactyl.py
+++ b/pterodactyl/pterodactyl.py
@@ -303,3 +303,12 @@ class Pterodactyl(commands.Cog):
             return await ctx.send(f"Leave message is currently set to: {message}")
         await config.leave_msg.set(message)
         await ctx.send(f"Leave message set to: {message}")
+
+    @pterodactyl_config.command(name = "ip")
+    async def pterodactyl_config_mask_ip(self, ctx: commands.Context, mask: bool = None) -> None:
+        """Mask the IP addresses of users in console messages."""
+        if mask is None:
+            mask = await config.mask_ip()
+            return await ctx.send(f"IP masking is currently set to {mask}")
+        await config.mask_ip.set(mask)
+        await ctx.send(f"IP masking set to {mask}")
diff --git a/pterodactyl/websocket.py b/pterodactyl/websocket.py
index 5f8e159..95625ae 100644
--- a/pterodactyl/websocket.py
+++ b/pterodactyl/websocket.py
@@ -47,6 +47,8 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
             if json.loads(message)['event'] == 'console output' and await config.console_channel() is not None:
                 if await config.current_status() in ('running', 'offline', ''):
                     content = remove_ansi_escape_codes(json.loads(message)['args'][0])
+                    if await config.mask_ip() is True:
+                        content = mask_ip(content)
 
                     channel = coginstance.bot.get_channel(await config.console_channel())
                     if channel is not None:
@@ -250,3 +252,10 @@ async def generate_achievement_embed(username: str, achievement: str, challenge:
         embed.set_author(name=username, icon_url='https://seafsh.cc/u/j3AzqQ.png')
     embed.timestamp = discord.utils.utcnow()
     return embed
+
+def mask_ip(string: str) -> str:
+    def mask_ip(match):
+        ip = match.group(0)
+        masked_ip = '.'.join('*' * len(octet) for octet in ip.split('.'))
+        return masked_ip
+    return re.sub(r'\b(?:\d{1,3}\.){3}\d{1,3}\b', mask_ip, string)