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

Make screesnhot command more efficient

Browse source
This commit is contained in:
nex 2023-01-09 14:36:32 +00:00
parent ae1bed5f66
commit e5d09ac385
2 changed files with 20 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
cogs/other.py
View file

@ -32,9 +32,8 @@ class OtherCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@staticmethod
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:
if not Path("/usr/bin/firefox").exists():
driver = "chrome"
@ -74,11 +73,14 @@ class OtherCog(commands.Cog):
driver = webdriver.Firefox(service=service, options=options)
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 ctx.edit(content=f"Screenshotting {friendly_url}... (66%)")
await _edit(content=f"Screenshotting {friendly_url}... (66%)")
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)
data = await asyncio.to_thread(driver.get_screenshot_as_png)
_io = io.BytesIO()
@ -370,21 +372,21 @@ class OtherCog(commands.Cog):
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:
for line in await blacklist.readlines():
if not line.strip():
continue
if re.match(line.strip(), url.netloc):
return False
return "Local blacklist"
# return await ctx.edit(content="That domain is blacklisted.")
return True
async def dns_check() -> Optional[bool]:
async def dns_check() -> Optional[bool | str]:
try:
for response in await asyncio.to_thread(dns.resolver.resolve, url.hostname, "A"):
if response.address == "0.0.0.0":
return False
return "DNS blacklist"
except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
return
else:
@ -397,19 +399,25 @@ class OtherCog(commands.Cog):
],
return_when=asyncio.FIRST_COMPLETED,
)
done = done or pending
done = done.pop()
if done.result() is not True:
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})"
)
await asyncio.sleep(1)
await ctx.edit(content=f"Preparing to screenshot {friendly_url}... (16%)")
okay = await pending.pop()
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})"
)
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)

4
main.py
View file

@ -47,7 +47,7 @@ async def on_connect():
@bot.listen("on_application_command_error")
async def on_application_command_error(ctx: discord.ApplicationContext, error: Exception):
await ctx.respond("Application Command Error: `%r`" % error)
raise error from None
raise 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):
return
await ctx.reply("Command Error: `%r`" % error)
raise error from None
raise error
@bot.listen("on_application_command")