nex/college-bot-v1
1
0
Fork 0
mirror of https://github.com/nexy7574/LCC-bot.git synced 2024-09-19 18:16:34 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix screenshot command & improve

Browse source
This commit is contained in:
nex 2023-01-15 19:39:07 +00:00
parent c3f78c85c8
commit e6465eca94
1 changed files with 116 additions and 71 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

187
cogs/other.py
View file

@ -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)