mirror of
https://github.com/myned/modufur.git
synced 2024-11-01 13:02:38 +00:00
132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
import urlextract
|
|
import hikari
|
|
import lightbulb
|
|
import pysaucenao
|
|
|
|
from tools import components, scraper
|
|
|
|
|
|
plugin = lightbulb.Plugin("booru")
|
|
extractor = urlextract.URLExtract()
|
|
|
|
|
|
@plugin.command
|
|
# @lightbulb.option('attachment', 'Attachment(s) to reverse')
|
|
@lightbulb.option("url", "URL(s) to reverse, separated by space")
|
|
@lightbulb.command("reverse", "Reverse image search using SauceNAO & Kheina", ephemeral=True)
|
|
@lightbulb.implements(lightbulb.SlashCommand, lightbulb.MessageCommand)
|
|
async def reverse(context):
|
|
match context:
|
|
case lightbulb.SlashContext():
|
|
urls = extractor.find_urls(context.options.url or "", only_unique=True, with_schema_only=True)
|
|
|
|
if not urls:
|
|
await context.respond("***Invalid URL(s)***")
|
|
return
|
|
|
|
await _reverse(context, urls)
|
|
case lightbulb.MessageContext():
|
|
urls = extractor.find_urls(context.options.target.content or "", only_unique=True, with_schema_only=True)
|
|
urls += [attachment.url for attachment in context.options.target.attachments if attachment.url not in urls]
|
|
|
|
if not urls:
|
|
await context.respond("***No images found***")
|
|
return
|
|
|
|
selector = None
|
|
|
|
if len(urls) > 1:
|
|
selector = components.Selector(
|
|
pages=[
|
|
f"**Select potential images to search: `{urls.index(url) + 1}/{len(urls)}`**\n{url}"
|
|
for url in urls
|
|
],
|
|
buttons=[components.Back(), components.Forward(), components.Select(), components.Confirm()],
|
|
urls=urls,
|
|
)
|
|
|
|
await selector.send(context.interaction, ephemeral=True)
|
|
await selector.wait()
|
|
|
|
if selector.timed_out:
|
|
await context.interaction.edit_initial_response("***Timed out***", components=None)
|
|
return
|
|
|
|
urls = selector.selected
|
|
|
|
await _reverse(context, urls, selector=selector)
|
|
|
|
|
|
@reverse.set_error_handler()
|
|
async def on_reverse_error(event):
|
|
error = None
|
|
|
|
match event.exception.__cause__:
|
|
case pysaucenao.ShortLimitReachedException():
|
|
error = "***API limit reached. Please try again in a minute***"
|
|
case pysaucenao.DailyLimitReachedException():
|
|
error = "***Daily API limit reached. Please try again tomorrow***"
|
|
case pysaucenao.FileSizeLimitException() as url:
|
|
error = f"***Image file size too large***\n{url}"
|
|
case pysaucenao.ImageSizeException() as url:
|
|
error = f"***Image resolution too small***\n{url}"
|
|
case pysaucenao.InvalidImageException() as url:
|
|
error = f"***Invalid image***\n{url}"
|
|
case pysaucenao.UnknownStatusCodeException():
|
|
error = "***An unknown SauceNAO error has occurred. The service may be down***"
|
|
|
|
if error:
|
|
try:
|
|
await event.context.respond(error)
|
|
except:
|
|
await event.context.interaction.edit_initial_response(error, components=None)
|
|
|
|
return True
|
|
|
|
|
|
async def _reverse(context, urls, *, selector=None):
|
|
if not selector:
|
|
await context.respond(hikari.ResponseType.DEFERRED_MESSAGE_CREATE)
|
|
|
|
matches = await scraper.reverse(urls)
|
|
|
|
if not matches:
|
|
if selector:
|
|
await context.interaction.edit_initial_response("***No matches found***", components=None)
|
|
else:
|
|
await context.respond("***No matches found***")
|
|
return
|
|
|
|
pages = [
|
|
(
|
|
hikari.Embed(
|
|
title=match["artist"],
|
|
url=match["url"],
|
|
color=context.get_guild().get_my_member().get_top_role().color if context.get_guild() else "#1a1a1a",
|
|
)
|
|
.set_author(name=f"{match['similarity']}% Match")
|
|
.set_image(match["thumbnail"])
|
|
.set_footer(match["source"])
|
|
)
|
|
if match
|
|
else f"***No match found***\n{urls[index]}"
|
|
for index, match in enumerate(matches)
|
|
]
|
|
|
|
if len(pages) > 1:
|
|
selector = components.Selector(pages=pages, buttons=[components.Back(), components.Forward()], timeout=900)
|
|
|
|
await selector.send_edit(context.interaction)
|
|
else:
|
|
if selector:
|
|
await context.interaction.edit_initial_response(content=None, embed=pages[0], components=None)
|
|
else:
|
|
await context.respond(pages[0])
|
|
|
|
|
|
def load(bot):
|
|
bot.add_plugin(plugin)
|
|
|
|
|
|
def unload(bot):
|
|
bot.remove_plugin(plugin)
|