jarvis-bot/jarvis/cogs/twitter.py
2022-02-04 10:23:07 -07:00

221 lines
7.8 KiB
Python

"""J.A.R.V.I.S. Twitter Cog."""
import asyncio
import tweepy
from bson import ObjectId
from dis_snek import InteractionContext, Permissions, Scale, Snake
from dis_snek.models.discord.channel import GuildText
from dis_snek.models.discord.components import ActionRow, Select, SelectOption
from dis_snek.models.snek.application_commands import (
OptionTypes,
SlashCommandChoice,
slash_command,
slash_option,
)
from dis_snek.models.snek.command import check
from jarvis.config import get_config
from jarvis.db.models import Twitter
from jarvis.utils.permissions import admin_or_permissions
class TwitterCog(Scale):
"""J.A.R.V.I.S. Twitter Cog."""
def __init__(self, bot: Snake):
self.bot = bot
config = get_config()
auth = tweepy.AppAuthHandler(
config.twitter["consumer_key"], config.twitter["consumer_secret"]
)
self.api = tweepy.API(auth)
self._guild_cache = {}
self._channel_cache = {}
@slash_command(name="twitter", sub_cmd_name="follow", description="Follow a Twitter acount")
@slash_option(
name="handle", description="Twitter account", opt_type=OptionTypes.STRING, required=True
)
@slash_option(
name="channel",
description="Channel to post tweets to",
opt_type=OptionTypes.CHANNEL,
required=True,
)
@slash_option(
name="retweets",
description="Mirror re-tweets?",
opt_type=OptionTypes.STRING,
required=False,
choices=[
SlashCommandChoice(name="Yes", value="Yes"),
SlashCommandChoice(name="No", value="No"),
],
)
@check(admin_or_permissions(Permissions.MANAGE_GUILD))
async def _twitter_follow(
self, ctx: InteractionContext, handle: str, channel: GuildText, retweets: str = "Yes"
) -> None:
handle = handle.lower()
retweets = retweets == "Yes"
if len(handle) > 15:
await ctx.send("Invalid Twitter handle", ephemeral=True)
return
if not isinstance(channel, GuildText):
await ctx.send("Channel must be a text channel", ephemeral=True)
return
try:
latest_tweet = await asyncio.to_thread(self.api.user_timeline, screen_name=handle)[0]
except Exception:
await ctx.send(
"Unable to get user timeline. Are you sure the handle is correct?", ephemeral=True
)
return
count = Twitter.objects(guild=ctx.guild.id).count()
if count >= 12:
await ctx.send("Cannot follow more than 12 Twitter accounts", ephemeral=True)
return
exists = Twitter.objects(handle=handle, guild=ctx.guild.id)
if exists:
await ctx.send("Twitter handle already being followed in this guild", ephemeral=True)
return
t = Twitter(
handle=handle,
guild=ctx.guild.id,
channel=channel.id,
admin=ctx.author.id,
last_tweet=latest_tweet.id,
retweets=retweets,
)
t.save()
await ctx.send(f"Now following `@{handle}` in {channel.mention}")
@slash_command(name="twitter", sub_cmd_name="unfollow", description="Unfollow Twitter accounts")
@check(admin_or_permissions(Permissions.MANAGE_GUILD))
async def _twitter_unfollow(self, ctx: InteractionContext) -> None:
twitters = Twitter.objects(guild=ctx.guild.id)
if not twitters:
await ctx.send("You need to follow a Twitter account first", ephemeral=True)
return
options = []
handlemap = {str(x.id): x.handle for x in twitters}
for twitter in twitters:
option = SelectOption(label=twitter.handle, value=str(twitter.id))
options.append(option)
select = Select(
options=options, custom_id="to_delete", min_values=1, max_values=len(twitters)
)
components = [ActionRow(select)]
block = "\n".join(x.handle for x in twitters)
message = await ctx.send(
content=f"You are following the following accounts:\n```\n{block}\n```\n\n"
"Please choose accounts to unfollow",
components=components,
)
try:
context = await self.bot.wait_for_component(
check=lambda x: ctx.author.id == x.author.id,
messages=message,
timeout=60 * 5,
)
for to_delete in context.context.values:
_ = Twitter.objects(guild=ctx.guild.id, id=ObjectId(to_delete)).delete()
for row in components:
for component in row.components:
component.disabled = True
block = "\n".join(handlemap[x] for x in context.context.values)
await context.context.edit_origin(
content=f"Unfollowed the following:\n```\n{block}\n```", components=components
)
except asyncio.TimeoutError:
for row in components:
for component in row.components:
component.disabled = True
await message.edit(components=components)
@slash_command(
name="twitter", sub_cmd_name="retweets", description="Modify followed Twitter accounts"
)
@slash_option(
name="retweets",
description="Mirror re-tweets?",
opt_type=OptionTypes.STRING,
required=False,
choices=[
SlashCommandChoice(name="Yes", value="Yes"),
SlashCommandChoice(name="No", value="No"),
],
)
@check(admin_or_permissions(Permissions.MANAGE_GUILD))
async def _twitter_modify(self, ctx: InteractionContext, retweets: str) -> None:
retweets = retweets == "Yes"
twitters = Twitter.objects(guild=ctx.guild.id)
if not twitters:
await ctx.send("You need to follow a Twitter account first", ephemeral=True)
return
options = []
for twitter in twitters:
option = SelectOption(label=twitter.handle, value=str(twitter.id))
options.append(option)
select = Select(
options=options, custom_id="to_update", min_values=1, max_values=len(twitters)
)
components = [ActionRow(select)]
block = "\n".join(x.handle for x in twitters)
message = await ctx.send(
content=f"You are following the following accounts:\n```\n{block}\n```\n\n"
f"Please choose which accounts to {'un' if not retweets else ''}follow retweets from",
components=components,
)
try:
context = await self.bot.wait_for_component(
check=lambda x: ctx.author.id == x.author.id,
messages=message,
timeout=60 * 5,
)
handlemap = {str(x.id): x.handle for x in twitters}
for to_update in context.context.values:
t = Twitter.objects(guild=ctx.guild.id, id=ObjectId(to_update)).first()
t.retweets = retweets
t.save()
for row in components:
for component in row.components:
component.disabled = True
block = "\n".join(handlemap[x] for x in context.context.values)
await context.context.edit_origin(
content=(
f"{'Unfollowed' if not retweets else 'Followed'} "
"retweets from the following:"
f"\n```\n{block}\n```"
),
components=components,
)
except asyncio.TimeoutError:
for row in components:
for component in row.components:
component.disabled = True
await message.edit(components=components)
def setup(bot: Snake) -> None:
"""Add TwitterCog to J.A.R.V.I.S."""
TwitterCog(bot)