Add Custom Tools to GUI (spy.pet-like Feature WIP)

bot/bot.py

@@ -123,6 +123,7 @@ async def on_message(self, message):
                 await self.add_cog(ScriptCog(self))
 
                 console.print_info(f"Loaded script: {script_name}")
+                self.controller.add_startup_script(script)
 
             except Exception as e:
                 console.print_error(f"Error loading script: {script_name} - {e}")
@@ -224,7 +225,8 @@ async def on_message_delete(self, message):
         if message.author.id == self.user.id:
             return
         delete_time = time.time()
-        self.controller.gui.home_page.add_discord_log(message.author, message, delete_time)
+        if self.controller.gui and hasattr(self.controller.gui, 'tools_page'):
+            self.controller.gui.tools_page.message_logger_page.add_discord_log(message.author, message, delete_time)
 
     def run_bot(self):
         try:

bot/commands/util.py

@@ -350,7 +350,7 @@ async def spypet(self, ctx, member_id: int):
         mutual_guilds = [guild for guild in self.bot.guilds if guild.get_member(member_id)]
         data = {}
         tasks = []
-        sem = asyncio.Semaphore(15)
+        sem = asyncio.Semaphore(10)
         stop_event = asyncio.Event()
         last_saved_count = 0
 
@@ -363,6 +363,7 @@ def _save_data():
             if current_count == last_saved_count:
                 return
             last_saved_count = current_count
+
             with open(files.get_application_support() + "/data/spypet.json", "w") as f:
                 json.dump(data, f, indent=4)
             console.print_info(f"Auto-saved {current_count} messages.")
@@ -399,15 +400,18 @@ async def _fetch_context_channel(channel):
                 latest_msg = [msg async for msg in channel.history(limit=1)][0]
                 context = await self.bot.get_context(latest_msg)
 
+                console.success(f"Got context for {channel.guild.name} - {channel.name}")
+
                 return context.channel
             except Exception as e:
                 if "429" in str(e):
+                    console.error(f"Rate limited while fetching context for {channel.guild.name} - {channel.name}")
                     await asyncio.sleep(5)
                     return await _fetch_context_channel(channel)
 
                 return None
 
-        async def _get_messages(channel, delay=0.25):
+        async def _get_messages(channel, delay=0.25, oldest_first=False):
             async with sem:
                 try:
                     await asyncio.sleep(delay)
@@ -418,7 +422,8 @@ async def _get_messages(channel, delay=0.25):
                     messages = []
 
                     try:
-                        async for msg in channel.history(limit=999999999, oldest_first=False):
+                        async for msg in channel.history(limit=None, oldest_first=oldest_first):
+                            print(channel.name, msg.author.id, msg.content)
                             if msg.author.id == member_id:
                                 if len(msg.attachments) > 0:
                                     attachments = "\n".join([f"Attachment: {attachment.url}" for attachment in msg.attachments])
@@ -427,13 +432,20 @@ async def _get_messages(channel, delay=0.25):
                                     msg_string = f"[{msg.created_at.strftime('%Y-%m-%d %H:%M:%S')}] {msg.content}"
                                 messages.append(msg_string)
                                 _add_message(guild, channel, msg_string)
+                                console.print_info(f"Found message in {channel.guild.name} - {channel.name}")
 
                         if len(messages) > 0: 
                             console.print_success(f"Found messages in {channel.guild.name} - {channel.name}")
                         else:
                             console.print_error(f"Found no messages in {channel.guild.name} - {channel.name}")
-                    except:
-                        console.print_error(f"Failed to fetch messages in {channel.guild.name} - {channel.name}")
+                    except Exception as e:
+                        if "429" in str(e).lower():
+                            console.print_error("Rate limited! Waiting for 5 seconds...")
+                            await asyncio.sleep(5)
+                            return await _get_messages(channel, delay)
+                        else:
+                            console.print_error(f"Error in {channel.guild.name} - {channel.name}: {e}")
+                            return
 
                     _save_data()
                 except asyncio.CancelledError:
@@ -446,14 +458,33 @@ async def _get_messages(channel, delay=0.25):
                 finally:
                     _save_data()
 
-        delay = 0.5
+        async def _attempt_scrape(guild, delay):
+            tasks = []
+            console.info(f"Attempting to scrape {guild.name} - {guild.id}")
+
+            for channel in guild.channels:
+                if isinstance(channel, discord.TextChannel) and _get_permissions(channel):
+                    tasks.append(_get_messages(channel, delay, oldest_first=True))
+                    delay += 2
+
+            if len(tasks) == 0:
+                console.print_error(f"No valid channels in {guild.name} - {guild.id}")
+                return
+
+            try:
+                await asyncio.gather(*tasks)
+            except asyncio.CancelledError:
+                console.print_warning("Process was cancelled! Saving progress...")
+                stop_event.set()
+                _save_data()
+                raise
+
+        delay = 1
         autosave_task = asyncio.create_task(_autosave(5))
 
         for guild in mutual_guilds:
-            for channel in guild.text_channels:
-                if _get_permissions(channel):
-                    tasks.append(asyncio.create_task(_get_messages(channel, delay)))
-                    delay += 0.5
+            tasks.append(_attempt_scrape(guild, delay))
+            delay += 1.5
 
         await asyncio.gather(*tasks)
         stop_event.set()
@@ -464,7 +495,7 @@ async def _get_messages(channel, delay=0.25):
         console.print_info(f"Total messages: {_count_messages()}")
         console.print_info(f"Total guilds: {len(data)}")
         console.print_info(f"Total channels: {sum(len(channels) for channels in data.values())}")
-        await ctx.send(file=discord.File("data/spypet.json"), delete_after=self.cfg.get("message_settings")["auto_delete_delay"])
+        await ctx.send(file=discord.File(files.get_application_support() + "/data/spypet.json"), delete_after=self.cfg.get("message_settings")["auto_delete_delay"])
 
 def setup(bot):
     bot.add_cog(Util(bot))

bot/controller.py

@@ -15,7 +15,7 @@
 from bot.helpers import cmdhelper, imgembed
 import utils.webhook as webhook_client
 from gui.helpers.images import resize_and_sharpen
-from bot.helpers.spypet import Spypet
+from bot.tools import SpyPet
 
 if getattr(sys, 'frozen', False):
     os.chdir(os.path.dirname(sys.executable))
@@ -33,13 +33,42 @@ def __init__(self):
         self.bot_running = False
         self.startup_scripts = []
         self.presence = self.cfg.get_rich_presence()
-        self.spypet = Spypet()
+        self.spypet = SpyPet(self)
+
+    def start_spypet(self):
+        if not self.spypet.bot:
+            self.spypet.set_bot(self.bot)
+            console.success("SpyPet bot set successfully.")
+        else:
+            console.warning("SpyPet bot is already set.")
+
+        if not self.spypet.member_id:
+            console.error("SpyPet member ID is not set. Please set it in the settings.")
+            return
+
+        asyncio.run_coroutine_threadsafe(self.spypet.start(), self.loop)
+        console.success("SpyPet started successfully!")
+
+    def stop_spypet(self):
+        if self.spypet.running:
+            asyncio.run_coroutine_threadsafe(self.spypet.stop(), self.loop)
+            console.success("SpyPet stopped successfully!")
+        else:
+            console.warning("SpyPet is not running.")
+
+    def get_mutual_guilds_spypet(self):
+        if self.spypet.running:
+            return asyncio.run_coroutine_threadsafe(self.spypet.get_mutual_guilds(), self.loop).result()
+        else:
+            console.warning("SpyPet is not running. Cannot get mutual guilds.")
+            return []
 
     def add_startup_script(self, script):
         self.startup_scripts.append(script)
 
     def set_gui(self, gui):
         self.gui = gui
+        self.spypet.set_gui(gui.tools_page.spypet_page)
 
     def check_token(self):
         resp = requests.get("https://discord.com/api/v9/users/@me", headers={"Authorization": self.cfg.get("token")})
@@ -190,16 +219,19 @@ def get_user_from_id(self, user_id):
         return self.bot.get_user(user_id)
 
     def get_avatar_from_url(self, url, size=50, radius=5):
-        url = url.split("?")[0]
-        if url.endswith(".gif"):
-            url = url.replace(".gif", ".png")
-        response = requests.get(url)
-        if response.status_code == 200:
-            image = Image.open(BytesIO(response.content))
-            # image = image.resize((size, size))
-            image = resize_and_sharpen(image, (size, size))
-            image = imgembed.add_corners(image, radius)
-            return ImageTk.PhotoImage(image)
+        try:
+            url = url.split("?")[0]
+            if url.endswith(".gif"):
+                url = url.replace(".gif", ".png")
+            response = requests.get(url, timeout=10)
+            if response.status_code == 200:
+                image = Image.open(BytesIO(response.content))
+                # image = image.resize((size, size))
+                image = resize_and_sharpen(image, (size, size))
+                image = imgembed.add_corners(image, radius)
+                return ImageTk.PhotoImage(image)
+        except Exception as e:
+            print(f"Error processing avatar from URL {url}: {e}")
 
         return None
 
@@ -216,6 +248,18 @@ def get_avatar(self, size=50, radius=5):
     def set_prefix(self, prefix):
         self.bot.command_prefix = prefix
 
+
+    async def get_user_from_id_async(self, user_id):
+        print(f"[BotController] Getting user from ID: {user_id}")
+        try:
+            return await self.bot.fetch_user(user_id) if self.bot else None
+        except Exception as e:
+            console.print_error(f"Error getting user from ID {user_id}: {e}")
+            return None
+
+    def get_user_from_id(self, user_id):
+        return asyncio.run_coroutine_threadsafe(self.get_user_from_id_async(user_id), self.loop).result()
+
     get_user    = lambda self: self.bot.user if self.bot else None
     get_friends = lambda self: self.bot.friends if self.bot else None
     get_guilds  = lambda self: self.bot.guilds if self.bot else None

bot/helpers/cmdhelper.py

@@ -100,6 +100,8 @@ async def rich_embed(ctx, embed):
     webhook_client.send(embed=embed.to_dict())
 
     async for message in webhook_channel.history(limit=1):
+        await message.ack()
+
         try:
             resp = requests.post(
                 f"https://discord.com/api/v9/channels/{ctx.channel.id}/messages", 

bot/helpers/imgembed.py

@@ -285,7 +285,10 @@ def draw(self):
     def save(self):
         path = f"embed-{random.randint(1000, 9999)}.png"  # comment this out if youre running the script directly
         # path = "embed.png" # uncomment this if youre running the script directly
-        self.draw().save(path)
+        final = self.draw()
+        final.thumbnail((self.width // 2, self.height // 2), Image.LANCZOS)
+        # final = final.resize((self.width // 2, self.height // 2), Image.LANCZOS)
+        final.save(path, optimize=True, quality=20)
         return path
 
 

bot/tools/__init__.py

@@ -0,0 +1 @@
+from .spypet import SpyPet