162 lines
5.0 KiB
Python
162 lines
5.0 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
from shlex import join
|
|
from tempfile import TemporaryDirectory
|
|
|
|
import discord
|
|
|
|
import discord_tools.message
|
|
import discord_tools.pixiv_link_extractor
|
|
import pixiv.pixiv_api as pixiv_api
|
|
|
|
logging.basicConfig(
|
|
format="[%(asctime)s] [%(levelname)8s] %(name)s: %(message)s", level=logging.INFO
|
|
)
|
|
|
|
|
|
REFRESH_TOKEN = os.getenv("PIXIV_REFRESH_TOKEN")
|
|
if REFRESH_TOKEN is None:
|
|
logging.fatal("Pixiv refresh token is missing, please set PIXIV_REFRESH_TOKEN")
|
|
exit(1)
|
|
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
|
|
if DISCORD_TOKEN is None:
|
|
logging.fatal("Discord token is missing, please set DISCORD_TOKEN")
|
|
exit(1)
|
|
|
|
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
|
|
|
|
class MyClient(discord.Client):
|
|
def __init__(self, *, intents: discord.Intents):
|
|
super().__init__(intents=intents)
|
|
self.tree = discord.app_commands.CommandTree(self)
|
|
|
|
async def setup_hook(self):
|
|
await self.tree.sync()
|
|
|
|
|
|
client = MyClient(intents=intents)
|
|
|
|
|
|
@client.event
|
|
async def on_ready():
|
|
logging.info(f"We have logged in as {client.user}")
|
|
|
|
|
|
@client.tree.command()
|
|
@discord.app_commands.describe(
|
|
url="Pixiv art url",
|
|
)
|
|
async def pixivprev(interaction: discord.Interaction, url: str):
|
|
"""Create pixiv preview"""
|
|
(result, match) = discord_tools.pixiv_link_extractor.check_pixiv_url(url)
|
|
if result and match is not None:
|
|
await interaction.response.defer()
|
|
with TemporaryDirectory() as tmp_dir:
|
|
loop = asyncio.get_running_loop()
|
|
data = await loop.run_in_executor(
|
|
None,
|
|
pixiv_api.download_pixiv_images,
|
|
int(match.group("id")),
|
|
tmp_dir,
|
|
REFRESH_TOKEN,
|
|
)
|
|
if data is None:
|
|
await interaction.response.send_message(
|
|
"Fail to extract pixiv image", ephemeral=True
|
|
)
|
|
return
|
|
(title, author) = data
|
|
files = os.listdir(tmp_dir)
|
|
files.sort()
|
|
await discord_tools.message.send_command_reply(
|
|
interaction,
|
|
files,
|
|
tmp_dir,
|
|
True,
|
|
title=title,
|
|
art_author=author,
|
|
url=url,
|
|
author=interaction.user,
|
|
)
|
|
else:
|
|
await interaction.response.send_message("Invalid pixiv url", ephemeral=True)
|
|
|
|
|
|
@client.tree.context_menu(name="Create pixiv preview")
|
|
async def pixiv_preview(interaction: discord.Interaction, message: discord.Message):
|
|
(result, muted, match) = discord_tools.pixiv_link_extractor.should_process(message)
|
|
if result and match is not None:
|
|
await interaction.response.defer()
|
|
with TemporaryDirectory() as tmp_dir:
|
|
loop = asyncio.get_running_loop()
|
|
data = await loop.run_in_executor(
|
|
None,
|
|
pixiv_api.download_pixiv_images,
|
|
int(match.group("id")),
|
|
tmp_dir,
|
|
REFRESH_TOKEN,
|
|
)
|
|
if data is None:
|
|
await interaction.response.send_message(
|
|
"Fail to extract pixiv image", ephemeral=True
|
|
)
|
|
return
|
|
files = os.listdir(tmp_dir)
|
|
files.sort()
|
|
(title, author) = data
|
|
await discord_tools.message.send_command_reply(
|
|
interaction=interaction,
|
|
files=files,
|
|
tmp_dir=tmp_dir,
|
|
muted=muted,
|
|
title=title,
|
|
art_author=author,
|
|
url=match.group("url"),
|
|
author=message.author,
|
|
)
|
|
else:
|
|
await interaction.response.send_message(
|
|
"Can't find any valid pixiv link.", ephemeral=True
|
|
)
|
|
|
|
|
|
@client.event
|
|
async def on_guild_join(guild: discord.Guild):
|
|
logging.info(f"Joined new guild: {guild.name}")
|
|
|
|
|
|
@client.event
|
|
async def on_message(message: discord.Message):
|
|
if message.author == client.user:
|
|
return
|
|
(result, muted, match) = discord_tools.pixiv_link_extractor.should_process(message)
|
|
if result and match is not None:
|
|
with TemporaryDirectory() as tmp_dir:
|
|
async with message.channel.typing():
|
|
loop = asyncio.get_running_loop()
|
|
data = await loop.run_in_executor(
|
|
None,
|
|
pixiv_api.download_pixiv_images,
|
|
int(match.group("id")),
|
|
tmp_dir,
|
|
REFRESH_TOKEN,
|
|
)
|
|
if data is None:
|
|
return
|
|
(title, author) = data
|
|
files = os.listdir(tmp_dir)
|
|
files.sort()
|
|
# await send_message.send_message_with_embed(
|
|
# message, files, title, tmp_dir, match.group("url")
|
|
# )
|
|
await discord_tools.message.send_message(
|
|
message, files, tmp_dir, muted, title, author
|
|
)
|
|
|
|
|
|
client.run(DISCORD_TOKEN)
|