mirror of
https://github.com/nexy7574/LCC-bot.git
synced 2024-09-19 18:16:34 +01:00
Fix screenshot command & improve
This commit is contained in:
parent
c3f78c85c8
commit
e6465eca94
1 changed files with 116 additions and 71 deletions
|
@ -33,23 +33,34 @@ 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:
|
||||
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'
|
||||
"/usr/bin/firefox-esr",
|
||||
"/usr/bin/firefox",
|
||||
# '/usr/bin/geckodriver'
|
||||
],
|
||||
"chrome": [
|
||||
'/usr/bin/chromedriver',
|
||||
'/usr/bin/chromium',
|
||||
'/usr/bin/chrome'
|
||||
]
|
||||
# '/usr/bin/chromedriver',
|
||||
"/usr/bin/chromium",
|
||||
"/usr/bin/chrome",
|
||||
"/usr/bin/chrome-browser",
|
||||
],
|
||||
}
|
||||
selected_driver = driver
|
||||
driver = None
|
||||
driver_path = None
|
||||
arr = drivers.pop(selected_driver)
|
||||
for binary in arr:
|
||||
b = Path(binary).resolve()
|
||||
|
@ -72,19 +83,30 @@ class OtherCog(commands.Cog):
|
|||
break
|
||||
else:
|
||||
raise RuntimeError("No browser binary.")
|
||||
return driver, driver_path
|
||||
|
||||
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("--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)
|
||||
driver.set_window_size(window_height, window_width)
|
||||
else:
|
||||
options = FirefoxOptions()
|
||||
options.add_argument("--headless")
|
||||
|
@ -94,20 +116,34 @@ class OtherCog(commands.Cog):
|
|||
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.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,10 +472,7 @@ 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(
|
||||
|
@ -449,7 +492,9 @@ class OtherCog(commands.Cog):
|
|||
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)
|
||||
|
|
Loading…
Reference in a new issue