Add some more error handling to the screenshot command

.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>

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,10 +81,22 @@ 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()
+        try:
             service = await asyncio.to_thread(ChromeService)
             driver: webdriver.Chrome = await asyncio.to_thread(
                 webdriver.Chrome,
@@ -104,11 +114,18 @@ class ScreenshotCog(commands.Cog):
                     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()
+        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()
 

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

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",