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

98
.pre-commit-config.yaml
View file

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

134
jarvis_tasks/__init__.py
View file

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

14
jarvis_tasks/config.py
View file

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

94
jarvis_tasks/tasks/ban.py
View file

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

146
jarvis_tasks/tasks/reminder.py
View file

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

224
jarvis_tasks/tasks/twitter.py
View file

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

56
jarvis_tasks/tasks/warning.py
View file

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