2024-01-02 22:59:47 +00:00
|
|
|
import asyncio
|
2024-04-18 00:24:58 +01:00
|
|
|
import copy
|
2024-01-02 22:29:25 +00:00
|
|
|
import datetime
|
|
|
|
import io
|
|
|
|
import logging
|
2024-01-02 22:59:47 +00:00
|
|
|
import os
|
2024-01-02 22:29:25 +00:00
|
|
|
import tempfile
|
|
|
|
import time
|
2024-04-02 16:03:14 +01:00
|
|
|
import typing
|
2024-01-02 22:29:25 +00:00
|
|
|
from urllib.parse import urlparse
|
|
|
|
|
|
|
|
import discord
|
2024-01-06 22:32:43 +00:00
|
|
|
import selenium.common
|
2024-01-02 22:29:25 +00:00
|
|
|
from discord.ext import commands
|
2024-01-02 22:59:47 +00:00
|
|
|
from PIL import Image
|
2024-01-02 22:29:25 +00:00
|
|
|
from selenium import webdriver
|
|
|
|
from selenium.webdriver.chrome.options import Options as ChromeOptions
|
|
|
|
from selenium.webdriver.chrome.service import Service as ChromeService
|
|
|
|
|
2024-03-17 16:15:59 +00:00
|
|
|
from conf import CONFIG
|
|
|
|
|
2024-04-02 16:03:14 +01:00
|
|
|
RESOLUTIONS = {
|
|
|
|
"8k": "7680x4320",
|
|
|
|
"4k": "3840x2160",
|
|
|
|
"1440p": "2560x1440",
|
|
|
|
"1080p": "1920x1080",
|
|
|
|
"720p": "1280x720",
|
|
|
|
"480p": "854x480",
|
|
|
|
"360p": "640x360",
|
|
|
|
"240p": "426x240",
|
2024-04-16 00:46:26 +01:00
|
|
|
"144p": "256x144",
|
2024-04-02 16:03:14 +01:00
|
|
|
}
|
|
|
|
_RES_OPTION = discord.Option(
|
|
|
|
name="resolution",
|
|
|
|
description="The resolution of the browser, can be WxH, or a preset (e.g. 1080p).",
|
2024-04-16 00:46:26 +01:00
|
|
|
autocomplete=discord.utils.basic_autocomplete(tuple(RESOLUTIONS.keys())),
|
2024-04-02 16:03:14 +01:00
|
|
|
)
|
|
|
|
|
|
|
|
|
2024-01-02 22:29:25 +00:00
|
|
|
class ScreenshotCog(commands.Cog):
|
|
|
|
def __init__(self, bot: commands.Bot):
|
|
|
|
self.bot = bot
|
2024-01-04 15:43:40 +00:00
|
|
|
self.log = logging.getLogger("jimmy.cogs.screenshot")
|
2024-01-02 22:29:25 +00:00
|
|
|
|
2024-04-02 16:03:14 +01:00
|
|
|
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")
|
2024-01-02 22:59:47 +00:00
|
|
|
if os.getuid() == 0:
|
2024-04-02 16:03:14 +01:00
|
|
|
_chrome_options.add_argument("--no-sandbox")
|
|
|
|
_chrome_options.add_argument("--disable-dev-shm-usage")
|
|
|
|
_chrome_options.add_argument("--disable-setuid-sandbox")
|
2024-01-04 15:43:40 +00:00
|
|
|
self.log.warning("Running as root, disabling chrome sandbox.")
|
2024-04-02 16:03:14 +01:00
|
|
|
_chrome_options.add_argument(
|
2024-02-11 18:04:37 +00:00
|
|
|
"--user-agent=Mozi11a/5.0 "
|
|
|
|
"(X11; Linux x86_64) AppleWebKit/537.36 (KHTML, Like Gecko) Chrome/121.9.6167.160 Safari/537.36"
|
|
|
|
)
|
2024-01-02 22:29:25 +00:00
|
|
|
|
|
|
|
prefs = {
|
2024-01-02 22:59:47 +00:00
|
|
|
"download.open_pdf_in_system_reader": False,
|
|
|
|
"plugins.always_open_pdf_externally": False,
|
2024-01-02 22:29:25 +00:00
|
|
|
"download_restrictions": 3,
|
|
|
|
}
|
2024-04-16 00:46:26 +01:00
|
|
|
_chrome_options.add_experimental_option("prefs", prefs)
|
2024-04-02 16:03:14 +01:00
|
|
|
return _chrome_options
|
2024-01-02 22:29:25 +00:00
|
|
|
|
|
|
|
def compress_png(self, input_file: io.BytesIO) -> io.BytesIO:
|
|
|
|
img = Image.open(input_file)
|
|
|
|
img = img.convert("RGB")
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".webp") as file:
|
|
|
|
quality = 100
|
|
|
|
while quality > 0:
|
|
|
|
if quality == 100:
|
|
|
|
quality_r = 99
|
|
|
|
else:
|
|
|
|
quality_r = quality
|
|
|
|
self.log.debug("Compressing image with quality %d%%", quality_r)
|
|
|
|
img.save(file.name, "webp", quality=quality_r)
|
|
|
|
file.seek(0)
|
|
|
|
value = io.BytesIO(file.read())
|
2024-04-02 02:55:54 +01:00
|
|
|
if len(value.getvalue()) <= 8 * 1024 * 1024:
|
2024-01-02 22:29:25 +00:00
|
|
|
self.log.debug("%d%% was sufficient.", quality_r)
|
|
|
|
break
|
|
|
|
quality -= 15
|
|
|
|
else:
|
|
|
|
raise RuntimeError("Couldn't compress image.")
|
|
|
|
return value
|
|
|
|
|
2024-01-02 22:59:47 +00:00
|
|
|
# noinspection PyTypeChecker
|
2024-01-02 22:29:25 +00:00
|
|
|
@commands.slash_command()
|
|
|
|
async def screenshot(
|
2024-04-16 00:46:26 +01:00
|
|
|
self,
|
|
|
|
ctx: discord.ApplicationContext,
|
|
|
|
url: str,
|
|
|
|
load_timeout: int = 10,
|
|
|
|
render_timeout: int = None,
|
|
|
|
eager: bool = None,
|
|
|
|
resolution: typing.Annotated[str, _RES_OPTION] = "1440p",
|
|
|
|
use_proxy: bool = False,
|
2024-01-02 22:29:25 +00:00
|
|
|
):
|
|
|
|
"""Screenshots a webpage."""
|
|
|
|
await ctx.defer()
|
|
|
|
|
|
|
|
if eager is None:
|
|
|
|
eager = render_timeout is None
|
|
|
|
if render_timeout is None:
|
2024-01-15 11:02:31 +00:00
|
|
|
# render_timeout = 30 if eager else 10
|
|
|
|
render_timeout = 30
|
2024-01-02 22:29:25 +00:00
|
|
|
if not url.startswith("http"):
|
|
|
|
url = "https://" + url
|
2024-01-04 15:43:40 +00:00
|
|
|
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",
|
2024-04-16 00:46:26 +01:00
|
|
|
resolution,
|
2024-01-04 15:43:40 +00:00
|
|
|
)
|
2024-01-02 22:29:25 +00:00
|
|
|
parsed = urlparse(url)
|
|
|
|
await ctx.respond("Initialising...")
|
|
|
|
|
|
|
|
start_init = time.time()
|
2024-01-04 15:43:40 +00:00
|
|
|
try:
|
2024-04-02 16:03:14 +01:00
|
|
|
options = self.chrome_options()
|
|
|
|
if use_proxy is True and (server := CONFIG["screenshot"].get("proxy")):
|
|
|
|
use_proxy = True
|
2024-03-17 16:15:59 +00:00
|
|
|
options.add_argument("--proxy-server=" + server)
|
2024-04-02 16:03:14 +01:00
|
|
|
else:
|
|
|
|
use_proxy = False
|
2024-01-04 15:43:40 +00:00
|
|
|
service = await asyncio.to_thread(ChromeService)
|
2024-04-16 00:46:26 +01:00
|
|
|
driver: webdriver.Chrome = await asyncio.to_thread(webdriver.Chrome, service=service, options=options)
|
2024-01-04 15:43:40 +00:00
|
|
|
driver.set_page_load_timeout(load_timeout)
|
|
|
|
if resolution:
|
2024-04-02 16:03:14 +01:00
|
|
|
resolution = RESOLUTIONS.get(resolution.lower(), resolution)
|
2024-01-04 15:43:40 +00:00
|
|
|
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
|
2024-01-02 22:29:25 +00:00
|
|
|
end_init = time.time()
|
|
|
|
|
|
|
|
await ctx.edit(content=("Loading webpage..." if not eager else "Loading & screenshotting webpage..."))
|
|
|
|
start_request = time.time()
|
2024-01-04 15:43:40 +00:00
|
|
|
try:
|
|
|
|
await asyncio.to_thread(driver.get, url)
|
2024-01-06 22:32:43 +00:00
|
|
|
except selenium.common.WebDriverException as e:
|
2024-01-15 10:51:53 +00:00
|
|
|
await self.bot.loop.run_in_executor(None, driver.quit)
|
2024-01-06 22:32:43 +00:00
|
|
|
if "TimeoutException" in str(e):
|
|
|
|
return await ctx.respond("Timed out while loading webpage.")
|
|
|
|
else:
|
|
|
|
return await ctx.respond("Failed to load webpage:\n```\n%s\n```" % str(e.msg))
|
2024-01-04 15:43:40 +00:00
|
|
|
except Exception as e:
|
2024-01-15 10:51:53 +00:00
|
|
|
await self.bot.loop.run_in_executor(None, driver.quit)
|
2024-01-04 15:43:40 +00:00
|
|
|
await ctx.respond("Failed to get the webpage: " + str(e))
|
|
|
|
raise
|
2024-01-02 22:29:25 +00:00
|
|
|
end_request = time.time()
|
|
|
|
|
|
|
|
if not eager:
|
|
|
|
now = discord.utils.utcnow()
|
|
|
|
expires = now + datetime.timedelta(seconds=render_timeout)
|
|
|
|
await ctx.edit(content=f"Rendering (expires {discord.utils.format_dt(expires, 'R')})...")
|
|
|
|
start_wait = time.time()
|
|
|
|
await asyncio.sleep(render_timeout)
|
|
|
|
end_wait = time.time()
|
|
|
|
else:
|
|
|
|
start_wait = end_wait = 1
|
|
|
|
|
|
|
|
await ctx.edit(content="Saving screenshot...")
|
|
|
|
start_save = time.time()
|
|
|
|
ss = await asyncio.to_thread(driver.get_screenshot_as_png)
|
|
|
|
file = io.BytesIO()
|
|
|
|
await asyncio.to_thread(file.write, ss)
|
|
|
|
file.seek(0)
|
|
|
|
end_save = time.time()
|
|
|
|
|
2024-04-02 02:55:54 +01:00
|
|
|
if len(await asyncio.to_thread(file.getvalue)) > 8 * 1024 * 1024:
|
2024-03-18 23:57:13 +00:00
|
|
|
await ctx.edit(content="Compressing screenshot...")
|
2024-01-02 22:29:25 +00:00
|
|
|
start_compress = time.time()
|
|
|
|
file = await asyncio.to_thread(self.compress_png, file)
|
|
|
|
fn = "screenshot.webp"
|
|
|
|
end_compress = time.time()
|
|
|
|
else:
|
|
|
|
fn = "screenshot.png"
|
|
|
|
start_compress = end_compress = 1
|
|
|
|
|
|
|
|
await ctx.edit(content="Cleaning up...")
|
|
|
|
start_cleanup = time.time()
|
2024-01-15 10:51:53 +00:00
|
|
|
await self.bot.loop.run_in_executor(None, driver.quit)
|
2024-01-02 22:29:25 +00:00
|
|
|
end_cleanup = time.time()
|
|
|
|
|
|
|
|
screenshot_size_mb = round(len(await asyncio.to_thread(file.getvalue)) / 1024 / 1024, 2)
|
|
|
|
|
|
|
|
def seconds(start: float, end: float) -> float:
|
|
|
|
return round(end - start, 2)
|
|
|
|
|
|
|
|
embed = discord.Embed(
|
|
|
|
title=f"Screenshot of {parsed.hostname}",
|
|
|
|
description=f"Init time: {seconds(start_init, end_init)}s\n"
|
2024-04-16 00:46:26 +01:00
|
|
|
f"Request time: {seconds(start_request, end_request)}s\n"
|
|
|
|
f"Wait time: {seconds(start_wait, end_wait)}s\n"
|
|
|
|
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"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}",
|
2024-01-02 22:29:25 +00:00
|
|
|
colour=discord.Colour.dark_theme(),
|
2024-04-16 00:46:26 +01:00
|
|
|
timestamp=discord.utils.utcnow(),
|
2024-01-02 22:29:25 +00:00
|
|
|
)
|
|
|
|
embed.set_image(url="attachment://" + fn)
|
|
|
|
return await ctx.edit(content=None, embed=embed, file=discord.File(file, filename=fn))
|
|
|
|
|
|
|
|
|
|
|
|
def setup(bot):
|
|
|
|
bot.add_cog(ScreenshotCog(bot))
|