mirror of
https://github.com/nexy7574/LCC-bot.git
synced 2024-09-19 18:16:34 +01:00
Make screesnhot command more efficient
This commit is contained in:
parent
ae1bed5f66
commit
e5d09ac385
2 changed files with 20 additions and 12 deletions
|
@ -32,9 +32,8 @@ class OtherCog(commands.Cog):
|
||||||
def __init__(self, bot):
|
def __init__(self, bot):
|
||||||
self.bot = bot
|
self.bot = bot
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
async def screenshot_website(
|
async def screenshot_website(
|
||||||
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
|
||||||
) -> discord.File:
|
) -> discord.File:
|
||||||
if not Path("/usr/bin/firefox").exists():
|
if not Path("/usr/bin/firefox").exists():
|
||||||
driver = "chrome"
|
driver = "chrome"
|
||||||
|
@ -74,11 +73,14 @@ class OtherCog(commands.Cog):
|
||||||
driver = webdriver.Firefox(service=service, options=options)
|
driver = webdriver.Firefox(service=service, options=options)
|
||||||
friendly_url = textwrap.shorten(website, 100)
|
friendly_url = textwrap.shorten(website, 100)
|
||||||
|
|
||||||
await ctx.edit(content=f"Screenshotting {friendly_url}... (49%)")
|
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 asyncio.to_thread(driver.get, website)
|
||||||
await ctx.edit(content=f"Screenshotting {friendly_url}... (66%)")
|
await _edit(content=f"Screenshotting {friendly_url}... (66%)")
|
||||||
await asyncio.sleep(render_time)
|
await asyncio.sleep(render_time)
|
||||||
await ctx.edit(content="Screenshotting {friendly_url}... (83%)")
|
await _edit(content=f"Screenshotting {friendly_url}... (83%)")
|
||||||
domain = re.sub(r"https?://", "", website)
|
domain = re.sub(r"https?://", "", website)
|
||||||
data = await asyncio.to_thread(driver.get_screenshot_as_png)
|
data = await asyncio.to_thread(driver.get_screenshot_as_png)
|
||||||
_io = io.BytesIO()
|
_io = io.BytesIO()
|
||||||
|
@ -370,21 +372,21 @@ class OtherCog(commands.Cog):
|
||||||
content=f"Preparing to screenshot {friendly_url}... (0%)"
|
content=f"Preparing to screenshot {friendly_url}... (0%)"
|
||||||
)
|
)
|
||||||
|
|
||||||
async def blacklist_check() -> bool:
|
async def blacklist_check() -> bool | str:
|
||||||
async with aiofiles.open("domains.txt") as blacklist:
|
async with aiofiles.open("domains.txt") as blacklist:
|
||||||
for line in await blacklist.readlines():
|
for line in await blacklist.readlines():
|
||||||
if not line.strip():
|
if not line.strip():
|
||||||
continue
|
continue
|
||||||
if re.match(line.strip(), url.netloc):
|
if re.match(line.strip(), url.netloc):
|
||||||
return False
|
return "Local blacklist"
|
||||||
# return await ctx.edit(content="That domain is blacklisted.")
|
# return await ctx.edit(content="That domain is blacklisted.")
|
||||||
return True
|
return True
|
||||||
|
|
||||||
async def dns_check() -> Optional[bool]:
|
async def dns_check() -> Optional[bool | str]:
|
||||||
try:
|
try:
|
||||||
for response in await asyncio.to_thread(dns.resolver.resolve, url.hostname, "A"):
|
for response in await asyncio.to_thread(dns.resolver.resolve, url.hostname, "A"):
|
||||||
if response.address == "0.0.0.0":
|
if response.address == "0.0.0.0":
|
||||||
return False
|
return "DNS blacklist"
|
||||||
except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
|
except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
|
||||||
return
|
return
|
||||||
else:
|
else:
|
||||||
|
@ -397,19 +399,25 @@ class OtherCog(commands.Cog):
|
||||||
],
|
],
|
||||||
return_when=asyncio.FIRST_COMPLETED,
|
return_when=asyncio.FIRST_COMPLETED,
|
||||||
)
|
)
|
||||||
|
done = done or pending
|
||||||
done = done.pop()
|
done = done.pop()
|
||||||
if done.result() is not True:
|
result = await done
|
||||||
|
if result is not True:
|
||||||
return await ctx.edit(
|
return await ctx.edit(
|
||||||
content="That domain is blacklisted, doesn't exist, or there was no answer from the DNS server."
|
content="That domain is blacklisted, doesn't exist, or there was no answer from the DNS server."
|
||||||
|
f" ({result!r})"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
await asyncio.sleep(1)
|
||||||
await ctx.edit(content=f"Preparing to screenshot {friendly_url}... (16%)")
|
await ctx.edit(content=f"Preparing to screenshot {friendly_url}... (16%)")
|
||||||
okay = await pending.pop()
|
okay = await pending.pop()
|
||||||
if okay is not True:
|
if okay is not True:
|
||||||
return await ctx.edit(
|
return await ctx.edit(
|
||||||
content="That domain is blacklisted, doesn't exist, or there was no answer from the DNS server."
|
content="That domain is blacklisted, doesn't exist, or there was no answer from the DNS server."
|
||||||
|
f" ({result!r})"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
await asyncio.sleep(1)
|
||||||
await ctx.edit(content=f"Screenshotting {textwrap.shorten(url.geturl(), 100)}... (33%)")
|
await ctx.edit(content=f"Screenshotting {textwrap.shorten(url.geturl(), 100)}... (33%)")
|
||||||
try:
|
try:
|
||||||
screenshot = await self.screenshot_website(ctx, url.geturl(), browser, render_timeout)
|
screenshot = await self.screenshot_website(ctx, url.geturl(), browser, render_timeout)
|
||||||
|
|
4
main.py
4
main.py
|
@ -47,7 +47,7 @@ async def on_connect():
|
||||||
@bot.listen("on_application_command_error")
|
@bot.listen("on_application_command_error")
|
||||||
async def on_application_command_error(ctx: discord.ApplicationContext, error: Exception):
|
async def on_application_command_error(ctx: discord.ApplicationContext, error: Exception):
|
||||||
await ctx.respond("Application Command Error: `%r`" % error)
|
await ctx.respond("Application Command Error: `%r`" % error)
|
||||||
raise error from None
|
raise error
|
||||||
|
|
||||||
|
|
||||||
@bot.listen("on_command_error")
|
@bot.listen("on_command_error")
|
||||||
|
@ -55,7 +55,7 @@ async def on_command_error(ctx: commands.Context, error: Exception):
|
||||||
if isinstance(error, commands.CommandNotFound):
|
if isinstance(error, commands.CommandNotFound):
|
||||||
return
|
return
|
||||||
await ctx.reply("Command Error: `%r`" % error)
|
await ctx.reply("Command Error: `%r`" % error)
|
||||||
raise error from None
|
raise error
|
||||||
|
|
||||||
|
|
||||||
@bot.listen("on_application_command")
|
@bot.listen("on_application_command")
|
||||||
|
|
Loading…
Reference in a new issue