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 @@
-
+
+
+
+
@@ -68,6 +71,8 @@
+
+
@@ -93,7 +98,15 @@
1704236387023
-
+
+
+ 1704296742685
+
+
+
+ 1704296742686
+
+
@@ -104,6 +117,7 @@
-
+
+
\ 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",