129 lines
5.2 KiB
Python
129 lines
5.2 KiB
Python
import re
|
|
import traceback
|
|
from aiohttp import ClientSession
|
|
import discord
|
|
import dyphanbot.utils as utils
|
|
from dyphanbot import Plugin
|
|
|
|
from .engines import EngineError, EngineRateLimitError
|
|
from .engines.saucenao import SauceNao
|
|
|
|
from .utils import parse_image_url
|
|
|
|
|
|
class SaucePlz(Plugin):
|
|
"""
|
|
DyphanBot plugin for reverse image searching
|
|
"""
|
|
|
|
async def help(self, message, args):
|
|
prefix = self.get_local_prefix(message)
|
|
return {
|
|
"helptext": ("__*Experimental feature, expect bugs!*__\n"
|
|
"Finds the source of an image using reverse image search engines.\n"
|
|
"Currently, only searches [SauceNAO](https://saucenao.com) databases, "
|
|
"but will hopefully support more engines in the future."),
|
|
"shorthelp": "Finds the source of an image using SauceNAO.",
|
|
"color": discord.Colour(0x1),
|
|
"sections": [{
|
|
"name": "Usage",
|
|
"value": (
|
|
"> {0}saucepls\n> {0}saucepls `URL`\n"
|
|
"Post an image or video with the above command "
|
|
"or call it with a URL.\n"
|
|
"Also works with replies!\n"
|
|
"Alias: *sauceplz*"
|
|
).format(prefix),
|
|
"inline": False
|
|
}]
|
|
}
|
|
|
|
def start(self):
|
|
self._config_fn = "config.json"
|
|
self.config = self.load_json(self._config_fn, initial_data={
|
|
"saucenao": {
|
|
"api_key": "__SAUCENAO_API_KEY__"
|
|
}
|
|
}, save_json=self._save_config)
|
|
|
|
def _save_config(self, filename, data):
|
|
return self.dyphanbot.data.save_json(filename, data, indent=4)
|
|
|
|
async def lookup_sauce(self, message, url):
|
|
try:
|
|
sn_engine = SauceNao(config=self.config, loop=self.dyphanbot.loop)
|
|
result = await sn_engine.best_match(url, hide_nsfw=not message.channel.is_nsfw())
|
|
if not result:
|
|
return {"content": "Unable to find match."}
|
|
|
|
embed = result.generate_embed(requester=message.author)
|
|
return {"content": "Sauce Found?", "embed": embed}
|
|
except EngineRateLimitError as err:
|
|
traceback.print_exc()
|
|
return {"content": f"Ratelimited: {err}"}
|
|
except EngineError as err:
|
|
traceback.print_exc()
|
|
return {"content": f"Error: {err}"}
|
|
|
|
@Plugin.command
|
|
async def sauceplz(self, client, message, args):
|
|
url = None
|
|
if len(args) > 0:
|
|
url = " ".join(args).strip("<>")
|
|
|
|
pre_text = ""
|
|
target_message = message
|
|
if len(message.attachments) <= 0 and not url:
|
|
# check if message is a reply
|
|
if message.reference:
|
|
msg_ref = message.reference
|
|
if not msg_ref.resolved:
|
|
try:
|
|
target_message = message.channel.fetch_message(msg_ref.message_id)
|
|
except Exception:
|
|
return await message.reply("Unable to retrieve referenced message.")
|
|
elif isinstance(msg_ref.resolved, discord.DeletedReferencedMessage):
|
|
return await message.reply("Referenced message was deleted.")
|
|
else:
|
|
target_message = msg_ref.resolved
|
|
urls = re.findall(r'(https?://\S+)', target_message.content)
|
|
if urls:
|
|
if len(urls) > 1:
|
|
pre_text += "Multiple URLs found in referenced message. Using the first one.\n"
|
|
url = urls[0]
|
|
if len(target_message.attachments) <= 0 and not url:
|
|
return await message.reply("No attachment or URL found in referenced message.")
|
|
else:
|
|
return await message.reply("No attachment or URL provided.")
|
|
|
|
if len(target_message.attachments) >= 1 and url is not None:
|
|
pre_text += "Both attachment and URL provided. Using URL.\n"
|
|
elif len(target_message.attachments) > 1:
|
|
pre_text += "Multiple attachments found. Using the first one.\n"
|
|
|
|
response_msg = None
|
|
|
|
if not url:
|
|
response_msg = await message.reply(f"{pre_text}*Getting attachment...*", mention_author=False)
|
|
image = target_message.attachments[0]
|
|
url = image.url
|
|
|
|
if response_msg:
|
|
await response_msg.edit(
|
|
allowed_mentions=discord.AllowedMentions(replied_user=False),
|
|
content=f"{pre_text}*Looking for the sauce...*")
|
|
else:
|
|
response_msg = await message.reply(f"{pre_text}*Looking for the sauce...*", mention_author=False)
|
|
|
|
async with ClientSession(loop=self.dyphanbot.loop) as session:
|
|
url = await parse_image_url(session, url)
|
|
|
|
results = await self.lookup_sauce(message, url=url)
|
|
await response_msg.edit(
|
|
allowed_mentions=discord.AllowedMentions(replied_user=True),
|
|
**results)
|
|
|
|
@Plugin.command
|
|
async def saucepls(self, client, message, args):
|
|
# alias to sauceplz
|
|
return await self.sauceplz(client, message, args) |