diff --git a/src/cogs/screenshot.py b/src/cogs/screenshot.py index c5b7965..204da13 100644 --- a/src/cogs/screenshot.py +++ b/src/cogs/screenshot.py @@ -6,6 +6,7 @@ import os import tempfile import time import copy +import typing from urllib.parse import urlparse import discord @@ -19,22 +20,41 @@ from selenium.webdriver.chrome.service import Service as ChromeService from conf import CONFIG +RESOLUTIONS = { + "8k": "7680x4320", + "4k": "3840x2160", + "1440p": "2560x1440", + "1080p": "1920x1080", + "720p": "1280x720", + "480p": "854x480", + "360p": "640x360", + "240p": "426x240", + "144p": "256x144" +} +_RES_OPTION = discord.Option( + name="resolution", + description="The resolution of the browser, can be WxH, or a preset (e.g. 1080p).", + autocomplete=discord.utils.basic_autocomplete(tuple(RESOLUTIONS.keys())) +) + + 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-extensions") - self.chrome_options.add_argument("--incognito") - self.chrome_options.add_argument("--remote-debugging-port=9222") + def chrome_options(self) -> ChromeOptions: + _chrome_options = ChromeOptions() + _chrome_options.add_argument("--headless") + _chrome_options.add_argument("--disable-extensions") + _chrome_options.add_argument("--incognito") + _chrome_options.add_argument("--remote-debugging-port=9222") if os.getuid() == 0: - self.chrome_options.add_argument("--no-sandbox") - self.chrome_options.add_argument("--disable-dev-shm-usage") - self.chrome_options.add_argument("--disable-setuid-sandbox") + _chrome_options.add_argument("--no-sandbox") + _chrome_options.add_argument("--disable-dev-shm-usage") + _chrome_options.add_argument("--disable-setuid-sandbox") self.log.warning("Running as root, disabling chrome sandbox.") - self.chrome_options.add_argument( + _chrome_options.add_argument( "--user-agent=Mozi11a/5.0 " "(X11; Linux x86_64) AppleWebKit/537.36 (KHTML, Like Gecko) Chrome/121.9.6167.160 Safari/537.36" ) @@ -44,9 +64,10 @@ class ScreenshotCog(commands.Cog): "plugins.always_open_pdf_externally": False, "download_restrictions": 3, } - self.chrome_options.add_experimental_option( + _chrome_options.add_experimental_option( "prefs", prefs ) + return _chrome_options def compress_png(self, input_file: io.BytesIO) -> io.BytesIO: img = Image.open(input_file) @@ -79,7 +100,10 @@ class ScreenshotCog(commands.Cog): load_timeout: int = 10, render_timeout: int = None, eager: bool = None, - resolution: str = "1920x1080", + resolution: typing.Annotated[ + str, + _RES_OPTION + ] = "1440p", use_proxy: bool = False ): """Screenshots a webpage.""" @@ -108,9 +132,12 @@ class ScreenshotCog(commands.Cog): start_init = time.time() try: - options = copy.copy(self.chrome_options) - if use_proxy and (server := CONFIG["screenshot"].get("proxy")): + options = self.chrome_options() + if use_proxy is True and (server := CONFIG["screenshot"].get("proxy")): + use_proxy = True options.add_argument("--proxy-server=" + server) + else: + use_proxy = False service = await asyncio.to_thread(ChromeService) driver: webdriver.Chrome = await asyncio.to_thread( webdriver.Chrome, @@ -119,18 +146,7 @@ class ScreenshotCog(commands.Cog): ) driver.set_page_load_timeout(load_timeout) if resolution: - common = { - "8k": "7680x4320", - "4k": "3840x2160", - "1440p": "2560x1440", - "1080p": "1920x1080", - "720p": "1280x720", - "480p": "854x480", - "360p": "640x360", - "240p": "426x240", - "144p": "256x144" - } - resolution = common.get(resolution.lower(), resolution) + resolution = RESOLUTIONS.get(resolution.lower(), resolution) try: width, height = map(int, resolution.split("x")) driver.set_window_size(width, height) @@ -207,7 +223,10 @@ class ScreenshotCog(commands.Cog): f"Save time: {seconds(start_save, end_save)}s\n" f"Compress time: {seconds(start_compress, end_compress)}s\n" f"Cleanup time: {seconds(start_cleanup, end_cleanup)}s\n" - f"Screenshot size: {screenshot_size_mb}MB\n", + f"Total time: {seconds(start_init, end_cleanup)}s\n" + f"Screenshot size: {screenshot_size_mb}MB\n" + f"Resolution: {resolution}" + f"Used proxy: {use_proxy}", colour=discord.Colour.dark_theme(), timestamp=discord.utils.utcnow() )