diff --git a/cogs/other.py b/cogs/other.py index 6f98b6b..5ec4990 100644 --- a/cogs/other.py +++ b/cogs/other.py @@ -33,81 +33,117 @@ class OtherCog(commands.Cog): self.bot = bot async def screenshot_website( - self, ctx: discord.ApplicationContext, website: str, driver: Literal["chrome", "firefox"], render_time: int = 10 + self, + ctx: discord.ApplicationContext, + website: str, + driver: Literal["chrome", "firefox"], + render_time: int = 10, + window_height: int = 1920, + window_width: int = 1080, + full_screenshot: bool = False, ) -> discord.File: - drivers = { - "firefox": [ - '/usr/bin/firefox-esr', - '/usr/bin/firefox', - '/usr/bin/geckodriver' - ], - "chrome": [ - '/usr/bin/chromedriver', - '/usr/bin/chromium', - '/usr/bin/chrome' - ] - } - selected_driver = driver - driver = None - driver_path = None - arr = drivers.pop(selected_driver) - for binary in arr: - b = Path(binary).resolve() - if not b.exists(): - continue - driver = selected_driver - driver_path = b - break - else: - for key, value in drivers.items(): - for binary in value: - b = Path(binary).resolve() - if not b.exists(): - continue - driver = key - driver_path = b - break - else: + async def _blocking(*args): + return await self.bot.loop.run_in_executor(None, *args) + + def find_driver(): + nonlocal driver + drivers = { + "firefox": [ + "/usr/bin/firefox-esr", + "/usr/bin/firefox", + # '/usr/bin/geckodriver' + ], + "chrome": [ + # '/usr/bin/chromedriver', + "/usr/bin/chromium", + "/usr/bin/chrome", + "/usr/bin/chrome-browser", + ], + } + selected_driver = driver + arr = drivers.pop(selected_driver) + for binary in arr: + b = Path(binary).resolve() + if not b.exists(): continue + driver = selected_driver + driver_path = b break else: - raise RuntimeError("No browser binary.") + for key, value in drivers.items(): + for binary in value: + b = Path(binary).resolve() + if not b.exists(): + continue + driver = key + driver_path = b + break + else: + continue + break + else: + raise RuntimeError("No browser binary.") + return driver, driver_path - if driver == "chrome": - options = ChromeOptions() - options.add_argument("--headless") - options.add_argument("--no-sandbox") - options.add_argument("--disable-dev-shm-usage") - options.add_argument("--disable-gpu") - options.add_argument("--window-size=1920x1080") - options.add_argument("--disable-extensions") - options.add_argument("--incognito") - options.binary_location = str(driver_path) - service = ChromeService("/usr/bin/chromedriver") - driver = webdriver.Chrome(service=service, options=options) - else: - options = FirefoxOptions() - options.add_argument("--headless") - options.add_argument("--private-window") - options.add_argument("--safe-mode") - options.add_argument("--new-instance") - options.binary_location = str(driver_path) - service = FirefoxService("/usr/bin/geckodriver") - driver = webdriver.Firefox(service=service, options=options) - friendly_url = textwrap.shorten(website, 100) + driver, driver_path = find_driver() + original_driver_name = driver + console.log( + "Using driver '{}' with binary '{}' to screenshot '{}', as requested by {}.".format( + driver, driver_path, website, ctx.user + ) + ) + + def _setup(): + nonlocal driver + if driver == "chrome": + options = ChromeOptions() + options.add_argument("--headless") + options.add_argument("--no-sandbox") + options.add_argument("--disable-dev-shm-usage") + options.add_argument("--disable-gpu") + options.add_argument("--disable-extensions") + options.add_argument("--incognito") + options.binary_location = str(driver_path) + service = ChromeService("/usr/bin/chromedriver") + driver = webdriver.Chrome(service=service, options=options) + driver.set_window_size(window_height, window_width) + else: + options = FirefoxOptions() + options.add_argument("--headless") + options.add_argument("--private-window") + options.add_argument("--safe-mode") + options.add_argument("--new-instance") + options.binary_location = str(driver_path) + service = FirefoxService("/usr/bin/geckodriver") + driver = webdriver.Firefox(service=service, options=options) + driver.set_window_size(window_height, window_width) + return driver, textwrap.shorten(website, 100) + + # Is it overkill to cast this to a thread? yes + # Do I give a flying fuck? kinda + # Why am I doing this? I suspect setup is causing a ~10-second block of the event loop + driver, friendly_url = await asyncio.to_thread(_setup) async def _edit(content: str): self.bot.loop.create_task(ctx.interaction.edit_original_response(content=content)) await _edit(content=f"Screenshotting {friendly_url}... (49%)") - await asyncio.to_thread(driver.get, website) + await _blocking(driver.set_page_load_timeout, render_time) + await _blocking(driver.get, website) await _edit(content=f"Screenshotting {friendly_url}... (66%)") await asyncio.sleep(render_time) await _edit(content=f"Screenshotting {friendly_url}... (83%)") domain = re.sub(r"https?://", "", website) - data = await asyncio.to_thread(driver.get_screenshot_as_png) + + screenshot_method = driver.get_screenshot_as_png + if full_screenshot and original_driver_name == "firefox": + screenshot_method = driver.get_full_page_screenshot_as_png + + data = await _blocking(screenshot_method) _io = io.BytesIO() - _io.write(data) + # Write the data async because HAHAHAHAHAHAHA + # We'll do it in the existing event loop though because less overhead + await _blocking(_io.write, data) _io.seek(0) driver.quit() return discord.File(_io, f"{domain}.png") @@ -377,7 +413,19 @@ class OtherCog(commands.Cog): ctx: discord.ApplicationContext, url: str, browser: discord.Option(str, description="Browser to use", choices=["chrome", "firefox"], default="firefox"), - render_timeout: int = 5, + render_timeout: discord.Option(int, name="render-timeout", description="Timeout for rendering", default=10), + window_height: discord.Option( + int, name="window-height", description="the height of the window in pixels", default=1920 + ), + window_width: discord.Option( + int, name="window-width", description="the width of the window in pixels", default=1080 + ), + capture_whole_page: discord.Option( + bool, + name="capture-full-page", + description="(firefox only) whether to capture the full page or just the viewport.", + default=False, + ) ): """Takes a screenshot of a URL""" await ctx.defer() @@ -391,9 +439,7 @@ class OtherCog(commands.Cog): url = urlparse(url) friendly_url = textwrap.shorten(url.geturl(), 100) - await ctx.edit( - content=f"Preparing to screenshot {friendly_url}... (0%)" - ) + await ctx.edit(content=f"Preparing to screenshot {friendly_url}... (0%)") async def blacklist_check() -> bool | str: async with aiofiles.open("domains.txt") as blacklist: @@ -426,15 +472,12 @@ class OtherCog(commands.Cog): try: done = done_tasks.pop() except KeyError: - return await ctx.respond( - "Something went wrong. Try again?\n" - f"`DONE:{done}|DONETASKS:{done_tasks}|PENDING:{pending}`" - ) + return await ctx.respond("Something went wrong. Try again?\n") result = await done if result is not True: return await ctx.edit( content="That domain is blacklisted, doesn't exist, or there was no answer from the DNS server." - f" ({result!r})" + f" ({result!r})" ) await asyncio.sleep(1) @@ -443,13 +486,15 @@ class OtherCog(commands.Cog): if okay is not True: return await ctx.edit( content="That domain is blacklisted, doesn't exist, or there was no answer from the DNS server." - f" ({result!r})" + f" ({result!r})" ) await asyncio.sleep(1) await ctx.edit(content=f"Screenshotting {textwrap.shorten(url.geturl(), 100)}... (33%)") try: - screenshot = await self.screenshot_website(ctx, url.geturl(), browser, render_timeout) + screenshot = await self.screenshot_website( + ctx, url.geturl(), browser, render_timeout, window_height, window_width, capture_whole_page + ) except Exception as e: console.print_exception() return await ctx.edit(content=f"Error: {e}", delete_after=30)