From 6fcd23ee35a8b16a63ea5befa1cc6dc0dc8b1f34 Mon Sep 17 00:00:00 2001
From: nexy7574 <git@nexy7574.co.uk>
Date: Thu, 4 Jan 2024 15:43:40 +0000
Subject: [PATCH] Add some more error handling to the screenshot command

---
 .idea/workspace.xml | 20 ++++++++++++---
 cogs/screenshot.py  | 62 +++++++++++++++++++++++++++++----------------
 conf.py             | 10 ++++++++
 main.py             |  6 +----
 4 files changed, 68 insertions(+), 30 deletions(-)
 create mode 100644 conf.py

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index 92ce689..259b868 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -4,8 +4,11 @@
     <option name="autoReloadType" value="SELECTIVE" />
   </component>
   <component name="ChangeListManager">
-    <list default="true" id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="Beta 1&#10;&#10;* Screenshot (unpolished, functional)&#10;* Dig&#10;* Traceroute&#10;* Whois&#10;* Ping&#10;* yt-dl&#10;* Docker file">
+    <list default="true" id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="Properly exit the webdriver">
+      <change afterPath="$PROJECT_DIR$/conf.py" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/cogs/screenshot.py" beforeDir="false" afterPath="$PROJECT_DIR$/cogs/screenshot.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/main.py" beforeDir="false" afterPath="$PROJECT_DIR$/main.py" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -68,6 +71,8 @@
       <workItem from="1704140889356" duration="166000" />
       <workItem from="1704141061129" duration="2495000" />
       <workItem from="1704224862150" duration="11897000" />
+      <workItem from="1704296563227" duration="780000" />
+      <workItem from="1704317484555" duration="1407000" />
     </task>
     <task id="LOCAL-00001" summary="Update gitignore">
       <option name="closed" value="true" />
@@ -93,7 +98,15 @@
       <option name="project" value="LOCAL" />
       <updated>1704236387023</updated>
     </task>
-    <option name="localTasksCounter" value="4" />
+    <task id="LOCAL-00004" summary="Update the docker image">
+      <option name="closed" value="true" />
+      <created>1704296742685</created>
+      <option name="number" value="00004" />
+      <option name="presentableId" value="LOCAL-00004" />
+      <option name="project" value="LOCAL" />
+      <updated>1704296742686</updated>
+    </task>
+    <option name="localTasksCounter" value="5" />
     <servers />
   </component>
   <component name="TypeScriptGeneratedFilesManager">
@@ -104,6 +117,7 @@
     <MESSAGE value="Update gitignore" />
     <MESSAGE value="Beta 1&#10;&#10;* Screenshot (unpolished, functional)&#10;* Dig&#10;* Traceroute&#10;* Whois&#10;* Ping&#10;* yt-dl" />
     <MESSAGE value="Beta 1&#10;&#10;* Screenshot (unpolished, functional)&#10;* Dig&#10;* Traceroute&#10;* Whois&#10;* Ping&#10;* yt-dl&#10;* Docker file" />
-    <option name="LAST_COMMIT_MESSAGE" value="Beta 1&#10;&#10;* Screenshot (unpolished, functional)&#10;* Dig&#10;* Traceroute&#10;* Whois&#10;* Ping&#10;* yt-dl&#10;* Docker file" />
+    <MESSAGE value="Update the docker image" />
+    <option name="LAST_COMMIT_MESSAGE" value="Update the docker image" />
   </component>
 </project>
\ No newline at end of file
diff --git a/cogs/screenshot.py b/cogs/screenshot.py
index 869e48e..4b11e65 100644
--- a/cogs/screenshot.py
+++ b/cogs/screenshot.py
@@ -18,15 +18,15 @@ from selenium.webdriver.chrome.service import Service as ChromeService
 class ScreenshotCog(commands.Cog):
     def __init__(self, bot: commands.Bot):
         self.bot = bot
+        self.log = logging.getLogger("jimmy.cogs.screenshot")
 
         self.chrome_options = ChromeOptions()
         self.chrome_options.add_argument("--headless")
-        self.chrome_options.add_argument("--disable-dev-shm-usage")
-        # self.chrome_options.add_argument("--disable-gpu")
         self.chrome_options.add_argument("--disable-extensions")
         self.chrome_options.add_argument("--incognito")
         if os.getuid() == 0:
             self.chrome_options.add_argument("--no-sandbox")
+            self.log.warning("Running as root, disabling chrome sandbox.")
 
         prefs = {
             "download.open_pdf_in_system_reader": False,
@@ -39,8 +39,6 @@ class ScreenshotCog(commands.Cog):
             "prefs", prefs
         )
 
-        self.log = logging.getLogger("jimmy.cogs.screenshot")
-
     def compress_png(self, input_file: io.BytesIO) -> io.BytesIO:
         img = Image.open(input_file)
         img = img.convert("RGB")
@@ -83,32 +81,51 @@ class ScreenshotCog(commands.Cog):
             render_timeout = 30 if eager else 10
         if not url.startswith("http"):
             url = "https://" + url
+        self.log.debug(
+            "User %s (%s) is attempting to screenshot %r with load timeout %d, render timeout %d, %s loading, and "
+            "a %s resolution.",
+            ctx.author,
+            ctx.author.id,
+            url,
+            load_timeout,
+            render_timeout,
+            "eager" if eager else "lazy",
+            resolution
+        )
         parsed = urlparse(url)
         await ctx.respond("Initialising...")
 
         start_init = time.time()
-        service = await asyncio.to_thread(ChromeService)
-        driver: webdriver.Chrome = await asyncio.to_thread(
-            webdriver.Chrome,
-            service=service,
-            options=self.chrome_options
-        )
-        driver.set_page_load_timeout(load_timeout)
-        if resolution:
-            try:
-                width, height = map(int, resolution.split("x"))
-                driver.set_window_size(width, height)
-                if height > 4320 or width > 7680:
-                    return await ctx.respond("Invalid resolution. Max resolution is 7680x4320 (8K).")
-            except ValueError:
-                return await ctx.respond("Invalid resolution. please provide width x height, e.g. 1920x1080")
-        if eager:
-            driver.implicitly_wait(render_timeout)
+        try:
+            service = await asyncio.to_thread(ChromeService)
+            driver: webdriver.Chrome = await asyncio.to_thread(
+                webdriver.Chrome,
+                service=service,
+                options=self.chrome_options
+            )
+            driver.set_page_load_timeout(load_timeout)
+            if resolution:
+                try:
+                    width, height = map(int, resolution.split("x"))
+                    driver.set_window_size(width, height)
+                    if height > 4320 or width > 7680:
+                        return await ctx.respond("Invalid resolution. Max resolution is 7680x4320 (8K).")
+                except ValueError:
+                    return await ctx.respond("Invalid resolution. please provide width x height, e.g. 1920x1080")
+            if eager:
+                driver.implicitly_wait(render_timeout)
+        except Exception as e:
+            await ctx.respond("Failed to initialise browser: " + str(e))
+            raise
         end_init = time.time()
 
         await ctx.edit(content=("Loading webpage..." if not eager else "Loading & screenshotting webpage..."))
         start_request = time.time()
-        await asyncio.to_thread(driver.get, url)
+        try:
+            await asyncio.to_thread(driver.get, url)
+        except Exception as e:
+            await ctx.respond("Failed to get the webpage: " + str(e))
+            raise
         end_request = time.time()
 
         if not eager:
@@ -140,6 +157,7 @@ class ScreenshotCog(commands.Cog):
 
         await ctx.edit(content="Cleaning up...")
         start_cleanup = time.time()
+        await asyncio.to_thread(driver.close)
         await asyncio.to_thread(driver.quit)
         end_cleanup = time.time()
 
diff --git a/conf.py b/conf.py
new file mode 100644
index 0000000..f82e55a
--- /dev/null
+++ b/conf.py
@@ -0,0 +1,10 @@
+import toml
+import logging
+from pathlib import Path
+
+try:
+    CONFIG = toml.load('config.toml')
+except FileNotFoundError:
+    cwd = Path.cwd()
+    logging.getLogger("jimmy.autoconf").critical("Unable to locate config.toml in %s.", cwd, exc_info=True)
+    raise
diff --git a/main.py b/main.py
index 1852a61..0b6f1cb 100644
--- a/main.py
+++ b/main.py
@@ -11,11 +11,7 @@ from rich.logging import RichHandler
 
 log = logging.getLogger("jimmy")
 
-try:
-    CONFIG = toml.load('config.toml')
-except FileNotFoundError:
-    log.critical("Unable to locate config.toml.", exc_info=True)
-    raise
+
 
 logging.basicConfig(
     format="%(asctime)s %(levelname)s %(name)s %(message)s",