jarvis/jarvis-tasks
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Changes

Browse source
This commit is contained in:
Zeva Rose 2022-03-10 17:22:24 -07:00
parent 933bff6711
commit 58c33d2077
8 changed files with 401 additions and 401 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
.flake8
View file

@ -1,18 +1,18 @@
[flake8] [flake8]
exclude = exclude =
run.py run.py
extend-ignore = extend-ignore =
Q0, E501, C812, E203, W503, # These default to arguing with Black. We might configure some of them eventually Q0, E501, C812, E203, W503, # These default to arguing with Black. We might configure some of them eventually
ANN001, # Ignore self and cls annotations ANN001, # Ignore self and cls annotations
ANN204, ANN206, # return annotations for special methods and class methods ANN204, ANN206, # return annotations for special methods and class methods
D105, D107, # Missing Docstrings in magic method and __init__ D105, D107, # Missing Docstrings in magic method and __init__
S311, # Standard pseudo-random generators are not suitable for security/cryptographic purposes. S311, # Standard pseudo-random generators are not suitable for security/cryptographic purposes.
D401, # First line should be in imperative mood; try rephrasing D401, # First line should be in imperative mood; try rephrasing
D400, # First line should end with a period D400, # First line should end with a period
D101, # Missing docstring in public class D101, # Missing docstring in public class
# Plugins we don't currently include: flake8-return # Plugins we don't currently include: flake8-return
R503, # missing explicit return at the end of function ableto return non-None value. R503, # missing explicit return at the end of function ableto return non-None value.
max-line-length=100 max-line-length=100

98
.pre-commit-config.yaml
View file

@ -1,49 +1,49 @@
repos: repos:
- repo: https://github.com/pre-commit/pre-commit-hooks - repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.1.0 rev: v4.1.0
hooks: hooks:
- id: check-toml - id: check-toml
- id: check-yaml - id: check-yaml
args: [--unsafe] args: [--unsafe]
- id: check-merge-conflict - id: check-merge-conflict
- id: requirements-txt-fixer - id: requirements-txt-fixer
- id: end-of-file-fixer - id: end-of-file-fixer
- id: debug-statements - id: debug-statements
language_version: python3.10 language_version: python3.10
- id: trailing-whitespace - id: trailing-whitespace
args: [--markdown-linebreak-ext=md] args: [--markdown-linebreak-ext=md]
- repo: https://github.com/pre-commit/pygrep-hooks - repo: https://github.com/pre-commit/pygrep-hooks
rev: v1.9.0 rev: v1.9.0
hooks: hooks:
- id: python-check-blanket-noqa - id: python-check-blanket-noqa
- repo: https://github.com/psf/black - repo: https://github.com/psf/black
rev: 22.1.0 rev: 22.1.0
hooks: hooks:
- id: black - id: black
args: [--line-length=100, --target-version=py310] args: [--line-length=100, --target-version=py310]
language_version: python3.10 language_version: python3.10
- repo: https://github.com/pre-commit/mirrors-isort - repo: https://github.com/pre-commit/mirrors-isort
rev: V5.10.1 rev: v5.10.1
hooks: hooks:
- id: isort - id: isort
args: ["--profile", "black"] args: ["--profile", "black"]
- repo: https://github.com/pycqa/flake8 - repo: https://github.com/pycqa/flake8
rev: 4.0.1 rev: 4.0.1
hooks: hooks:
- id: flake8 - id: flake8
additional_dependencies: additional_dependencies:
- flake8-annotations~=2.0 - flake8-annotations~=2.0
- flake8-bandit~=2.1 #- flake8-bandit~=2.1
- flake8-docstrings~=1.5 - flake8-docstrings~=1.5
- flake8-bugbear - flake8-bugbear
- flake8-comprehensions - flake8-comprehensions
- flake8-quotes - flake8-quotes
- flake8-raise - flake8-raise
- flake8-deprecated - flake8-deprecated
- flake8-print - flake8-print
- flake8-return - flake8-return
language_version: python3.10 language_version: python3.10

134
jarvis_tasks/__init__.py
View file

@ -1,67 +1,67 @@
"""JARVIS background tasks.""" """JARVIS background tasks."""
import asyncio import asyncio
from typing import Optional from typing import Optional
from dis_snek import Intents, Snake from dis_snek import Intents, Snake
from jarvis_core.db import connect from jarvis_core.db import connect
from jarvis_core.log import get_logger from jarvis_core.log import get_logger
from jarvis_tasks.config import TaskConfig from jarvis_tasks.config import TaskConfig
from jarvis_tasks.tasks import ban, reminder, twitter, warning from jarvis_tasks.tasks import ban, reminder, twitter, warning
logger = None logger = None
async def _start(config: Optional[str] = "config.yaml") -> None: async def _start(config: Optional[str] = "config.yaml") -> None:
""" """
Main start function. Main start function.
Args: Args:
config: Config path config: Config path
""" """
# Load config # Load config
config = TaskConfig.from_yaml(config) config = TaskConfig.from_yaml(config)
# Connect to database # Connect to database
testing = config.mongo["database"] != "jarvis" testing = config.mongo["database"] != "jarvis"
logger.debug(f"Connecting to database, testing={testing}") logger.debug(f"Connecting to database, testing={testing}")
connect(**config.mongo["connect"], testing=testing) connect(**config.mongo["connect"], testing=testing)
# Get event loop # Get event loop
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
# Login as bot # Login as bot
logger.debug("Logging in bot") logger.debug("Logging in bot")
intents = Intents.DEFAULT | Intents.GUILD_MEMBERS intents = Intents.DEFAULT | Intents.GUILD_MEMBERS
bot = Snake(intents=intents, loop=loop) bot = Snake(intents=intents, loop=loop)
await bot.login(config.token) await bot.login(config.token)
logger.info(f"Logged in as {bot.user.username}#{bot.user.discriminator}") logger.info(f"Logged in as {bot.user.username}#{bot.user.discriminator}")
# Start tasks # Start tasks
try: try:
logger.debug("Starting tasks") logger.debug("Starting tasks")
functions = [ban.unban, reminder.remind, twitter.twitter, warning.unwarn] functions = [ban.unban, reminder.remind, twitter.twitter, warning.unwarn]
tasks = [loop.create_task(f(bot, logger)) for f in functions] tasks = [loop.create_task(f(bot, logger)) for f in functions]
for task in tasks: for task in tasks:
await task await task
except KeyboardInterrupt: except KeyboardInterrupt:
for task in tasks: for task in tasks:
task.cancel() task.cancel()
def start(config: Optional[str] = "config.yaml") -> None: def start(config: Optional[str] = "config.yaml") -> None:
""" """
Start the background tasks. Start the background tasks.
Args: Args:
config: Config path config: Config path
""" """
global logger, debug global logger, debug
# Set log level # Set log level
_config = TaskConfig.from_yaml(config) _config = TaskConfig.from_yaml(config)
logger = get_logger(__name__) logger = get_logger(__name__)
logger.setLevel(_config.log_level) logger.setLevel(_config.log_level)
# Run the main tasks # Run the main tasks
logger.debug("Starting asyncio") logger.debug("Starting asyncio")
asyncio.run(_start(config)) asyncio.run(_start(config))

14
jarvis_tasks/config.py
View file

@ -1,7 +1,7 @@
"""Task config.""" """Task config."""
from jarvis_core.config import Config from jarvis_core.config import Config
class TaskConfig(Config): class TaskConfig(Config):
REQUIRED = ["token", "mongo", "twitter"] REQUIRED = ["token", "mongo", "twitter"]
OPTIONAL = {"log_level": "WARNING"} OPTIONAL = {"log_level": "WARNING"}

94
jarvis_tasks/tasks/ban.py
View file

@ -1,47 +1,47 @@
"""JARVIS ban tasks.""" """JARVIS ban tasks."""
import asyncio import asyncio
from datetime import datetime, timedelta from datetime import datetime, timedelta
from logging import Logger from logging import Logger
from dis_snek import Snake from dis_snek import Snake
from dis_snek.client.errors import NotFound from dis_snek.client.errors import NotFound
from jarvis_core.db import q from jarvis_core.db import q
from jarvis_core.db.models import Ban, Unban from jarvis_core.db.models import Ban, Unban
async def unban(bot: Snake, logger: Logger) -> None: async def unban(bot: Snake, logger: Logger) -> None:
""" """
Unban users when ban expires. Unban users when ban expires.
Args: Args:
bot: Snake instance bot: Snake instance
logger: Global logger logger: Global logger
""" """
while True: while True:
max_time = datetime.utcnow() + timedelta(minutes=10) max_time = datetime.utcnow() + timedelta(minutes=10)
bans = Ban.find(q(type="temp", active=True)) bans = Ban.find(q(type="temp", active=True))
async for ban in bans: async for ban in bans:
if ban.created_at + timedelta(hours=ban.duration) < max_time: if ban.created_at + timedelta(hours=ban.duration) < max_time:
guild = await bot.fetch_guild(ban.guild) guild = await bot.fetch_guild(ban.guild)
user = await bot.fetch_user(ban.user) user = await bot.fetch_user(ban.user)
if guild and user: if guild and user:
logger.debug(f"Unbanned user {user.id} from guild {guild.id}") logger.debug(f"Unbanned user {user.id} from guild {guild.id}")
try: try:
await guild.unban(user=user, reason="JARVIS tempban expired") await guild.unban(user=user, reason="JARVIS tempban expired")
except NotFound: except NotFound:
logger.debug(f"User {user.id} not banned from guild {guild.id}") logger.debug(f"User {user.id} not banned from guild {guild.id}")
ban.update(q(active=False)) ban.update(q(active=False))
await ban.commit() await ban.commit()
u = Unban( u = Unban(
user=user.id, user=user.id,
guild=guild.id, guild=guild.id,
username=user.username, username=user.username,
discrim=user.discriminator, discrim=user.discriminator,
admin=bot.user.id, admin=bot.user.id,
reason="Ban expired", reason="Ban expired",
) )
await u.commit() await u.commit()
# Check ever 10 minutes # Check ever 10 minutes
await asyncio.sleep(600) await asyncio.sleep(600)

146
jarvis_tasks/tasks/reminder.py
View file

@ -1,73 +1,73 @@
"""JARVIS reminders.""" """JARVIS reminders."""
import asyncio import asyncio
from datetime import datetime, timedelta from datetime import datetime, timedelta
from logging import Logger from logging import Logger
from dis_snek import Snake from dis_snek import Snake
from jarvis_core.db import q from jarvis_core.db import q
from jarvis_core.db.models import Reminder from jarvis_core.db.models import Reminder
from jarvis_core.util import build_embed from jarvis_core.util import build_embed
async def remind(bot: Snake, logger: Logger) -> None: async def remind(bot: Snake, logger: Logger) -> None:
""" """
Run reminders in the background. Run reminders in the background.
Args: Args:
bot: Snake instance bot: Snake instance
logger: Global logger logger: Global logger
""" """
while True: while True:
reminders = Reminder.find( reminders = Reminder.find(
q(remind_at__lte=datetime.utcnow() + timedelta(seconds=5), active=True) q(remind_at__lte=datetime.utcnow() + timedelta(seconds=5), active=True)
) )
async for reminder in reminders: async for reminder in reminders:
user = await bot.fetch_user(reminder.user) user = await bot.fetch_user(reminder.user)
if not user: if not user:
logger.warning(f"Failed to get user with ID {reminder.user}") logger.warning(f"Failed to get user with ID {reminder.user}")
await reminder.delete() await reminder.delete()
continue continue
embed = build_embed( embed = build_embed(
title="You have a reminder!", description=reminder.message, fields=[] title="You have a reminder!", description=reminder.message, fields=[]
) )
embed.set_author( embed.set_author(
name=user.username + "#" + user.discriminator, icon_url=user.avatar.url name=user.username + "#" + user.discriminator, icon_url=user.avatar.url
) )
embed.set_thumbnail(url=user.avatar.url) embed.set_thumbnail(url=user.avatar.url)
try: try:
await user.send(embed=embed) await user.send(embed=embed)
logger.info(f"Reminder {reminder.id} sent to user") logger.info(f"Reminder {reminder.id} sent to user")
await reminder.delete() await reminder.delete()
except Exception: except Exception:
logger.info("User has closed DMs") logger.info("User has closed DMs")
guild = await bot.fetch_guild(reminder.guild) guild = await bot.fetch_guild(reminder.guild)
member = await bot.fetch_member(user.id) member = await bot.fetch_member(user.id)
if not member: if not member:
logger.warning("User no longer member of origin guild, deleting reminder") logger.warning("User no longer member of origin guild, deleting reminder")
await reminder.delete() await reminder.delete()
continue continue
channel = await guild.fetch_channel(reminder.channel) if guild else None channel = await guild.fetch_channel(reminder.channel) if guild else None
if channel and not reminder.private: if channel and not reminder.private:
await channel.send(f"{member.mention}", embed=embed) await channel.send(f"{member.mention}", embed=embed)
logger.debug(f"Reminder {reminder.id} sent to origin channel") logger.debug(f"Reminder {reminder.id} sent to origin channel")
await reminder.delete() await reminder.delete()
elif channel: elif channel:
await channel.send( await channel.send(
f"{member.mention}, you had a private reminder set for now," f"{member.mention}, you had a private reminder set for now,"
" but I couldn't send it to you.\n" " but I couldn't send it to you.\n"
f"Use `/reminder fetch {str(reminder.id)}` to view" f"Use `/reminder fetch {str(reminder.id)}` to view"
) )
logger.info( logger.info(
f"Reminder {reminder.id} private, sent notification to origin channel" f"Reminder {reminder.id} private, sent notification to origin channel"
) )
reminder.update(q(active=False)) reminder.update(q(active=False))
await reminder.commit() await reminder.commit()
else: else:
logger.warning("No way to contact user, deleting reminder") logger.warning("No way to contact user, deleting reminder")
await reminder.delete() await reminder.delete()
# Check every 5 seconds # Check every 5 seconds
await asyncio.sleep(5) await asyncio.sleep(5)

224
jarvis_tasks/tasks/twitter.py
View file

@ -1,112 +1,112 @@
"""JARVIS Twitter sync.""" """JARVIS Twitter sync."""
import asyncio import asyncio
from datetime import datetime, timedelta from datetime import datetime, timedelta
from logging import Logger from logging import Logger
import tweepy import tweepy
from dis_snek import Snake from dis_snek import Snake
from dis_snek.models.discord.embed import EmbedAttachment from dis_snek.models.discord.embed import EmbedAttachment
from jarvis_core.db import q from jarvis_core.db import q
from jarvis_core.db.models import TwitterAccount, TwitterFollow from jarvis_core.db.models import TwitterAccount, TwitterFollow
from jarvis_core.util import build_embed from jarvis_core.util import build_embed
from jarvis_tasks.config import TaskConfig from jarvis_tasks.config import TaskConfig
config = TaskConfig.from_yaml() config = TaskConfig.from_yaml()
async def twitter(bot: Snake, logger: Logger) -> None: async def twitter(bot: Snake, logger: Logger) -> None:
""" """
Sync tweets in the background. Sync tweets in the background.
Args: Args:
bot: Snake instance bot: Snake instance
logger: Global logger logger: Global logger
""" """
auth = tweepy.AppAuthHandler(config.twitter["consumer_key"], config.twitter["consumer_secret"]) auth = tweepy.AppAuthHandler(config.twitter["consumer_key"], config.twitter["consumer_secret"])
api = tweepy.API(auth) api = tweepy.API(auth)
while True: while True:
accounts = TwitterAccount.find() accounts = TwitterAccount.find()
accounts_to_delete = [] accounts_to_delete = []
# Go through all actively followed accounts # Go through all actively followed accounts
async for account in accounts: async for account in accounts:
logger.debug(f"Checking account {account.handle}") logger.debug(f"Checking account {account.handle}")
# Check if account needs updated (handle changes) # Check if account needs updated (handle changes)
if account.last_sync + timedelta(hours=1) <= datetime.utcnow(): if account.last_sync + timedelta(hours=1) <= datetime.utcnow():
logger.debug(f"Account {account.handle} out of sync, updating") logger.debug(f"Account {account.handle} out of sync, updating")
user = api.get_user(user_id=account.twitter_id) user = api.get_user(user_id=account.twitter_id)
account.update(q(handle=user.screen_name, last_sync=datetime.utcnow())) account.update(q(handle=user.screen_name, last_sync=datetime.utcnow()))
# Get new tweets # Get new tweets
if tweets := api.user_timeline(user_id=account.twitter_id, since_id=account.last_tweet): if tweets := api.user_timeline(user_id=account.twitter_id, since_id=account.last_tweet):
logger.debug(f"{account.handle} has new tweets") logger.debug(f"{account.handle} has new tweets")
tweets = sorted(tweets, key=lambda x: x.id) tweets = sorted(tweets, key=lambda x: x.id)
follows = TwitterFollow.find(q(twitter_id=account.twitter_id)) follows = TwitterFollow.find(q(twitter_id=account.twitter_id))
follows_to_delete = [] follows_to_delete = []
num_follows = 0 num_follows = 0
# Go through follows and send tweet if necessary # Go through follows and send tweet if necessary
async for follow in follows: async for follow in follows:
num_follows += 1 num_follows += 1
guild = await bot.fetch_guild(follow.guild) guild = await bot.fetch_guild(follow.guild)
if not guild: if not guild:
logger.warning(f"Follow {follow.id}'s guild no longer exists, deleting") logger.warning(f"Follow {follow.id}'s guild no longer exists, deleting")
follows_to_delete.append(follow) follows_to_delete.append(follow)
continue continue
channel = await bot.fetch_channel(follow.channel) channel = await bot.fetch_channel(follow.channel)
if not channel: if not channel:
logger.warning(f"Follow {follow.id}'s channel no longer exists, deleting") logger.warning(f"Follow {follow.id}'s channel no longer exists, deleting")
follows_to_delete.append(follow) follows_to_delete.append(follow)
continue continue
for tweet in tweets: for tweet in tweets:
retweet = "retweeted_status" in tweet.__dict__ retweet = "retweeted_status" in tweet.__dict__
if retweet and not follow.retweets: if retweet and not follow.retweets:
continue continue
timestamp = int(tweet.created_at.timestamp()) timestamp = int(tweet.created_at.timestamp())
url = f"https://twitter.com/{account.handle}/status/{tweet.id}" url = f"https://twitter.com/{account.handle}/status/{tweet.id}"
mod = "re" if retweet else "" mod = "re" if retweet else ""
media = tweet.entities.get("media", None) media = tweet.entities.get("media", None)
photo = None photo = None
if media and media[0]["type"] in ["photo", "animated_gif"]: if media and media[0]["type"] in ["photo", "animated_gif"]:
photo = EmbedAttachment(url=media[0]["media_url_https"]) photo = EmbedAttachment(url=media[0]["media_url_https"])
embed = build_embed( embed = build_embed(
title="", title="",
description=(tweet.text + f"\n\n[View this tweet]({url})"), description=(tweet.text + f"\n\n[View this tweet]({url})"),
fields=[], fields=[],
color="#1DA1F2", color="#1DA1F2",
image=photo, image=photo,
) )
embed.set_author( embed.set_author(
name=account.handle, name=account.handle,
url=url, url=url,
icon_url=tweet.author.profile_image_url_https, icon_url=tweet.author.profile_image_url_https,
) )
embed.set_footer( embed.set_footer(
text="Twitter", text="Twitter",
icon_url="https://abs.twimg.com/icons/apple-touch-icon-192x192.png", icon_url="https://abs.twimg.com/icons/apple-touch-icon-192x192.png",
) )
await channel.send( await channel.send(
f"`@{account.handle}` {mod}tweeted this at <t:{timestamp}:f>" f"`@{account.handle}` {mod}tweeted this at <t:{timestamp}:f>"
) )
# Delete invalid follows # Delete invalid follows
for follow in follows_to_delete: for follow in follows_to_delete:
await follow.delete() await follow.delete()
if num_follows == 0: if num_follows == 0:
accounts_to_delete.append(account) accounts_to_delete.append(account)
else: else:
newest = tweets[-1] newest = tweets[-1]
account.update(q(last_tweet=newest.id)) account.update(q(last_tweet=newest.id))
await account.commit() await account.commit()
# Delete invalid accounts (no follows) # Delete invalid accounts (no follows)
for account in accounts_to_delete: for account in accounts_to_delete:
logger.info(f"Account {account.handle} has no followers, removing") logger.info(f"Account {account.handle} has no followers, removing")
await account.delete() await account.delete()
# Only check once a minute # Only check once a minute
await asyncio.sleep(60) await asyncio.sleep(60)

56
jarvis_tasks/tasks/warning.py
View file

@ -1,28 +1,28 @@
"""JARVIS warnings tasks.""" """JARVIS warnings tasks."""
import asyncio import asyncio
from datetime import datetime, timedelta from datetime import datetime, timedelta
from logging import Logger from logging import Logger
from dis_snek import Snake from dis_snek import Snake
from jarvis_core.db import q from jarvis_core.db import q
from jarvis_core.db.models import Warning from jarvis_core.db.models import Warning
async def unwarn(bot: Snake, logger: Logger) -> None: async def unwarn(bot: Snake, logger: Logger) -> None:
""" """
Deactivate warnings when they expire. Deactivate warnings when they expire.
Args: Args:
bot: Snake instance bot: Snake instance
logger: Global logger logger: Global logger
""" """
while True: while True:
warns = Warning.find(q(active=True)) warns = Warning.find(q(active=True))
async for warn in warns: async for warn in warns:
if warn.created_at + timedelta(hours=warn.duration) < datetime.utcnow(): if warn.created_at + timedelta(hours=warn.duration) < datetime.utcnow():
logger.debug(f"Deactivating warning {warn.id}") logger.debug(f"Deactivating warning {warn.id}")
warn.update(q(active=False)) warn.update(q(active=False))
await warn.commit() await warn.commit()
# Check every hour # Check every hour
await asyncio.sleep(3600) await asyncio.sleep(3600)