91 lines
2.2 KiB
Python
91 lines
2.2 KiB
Python
"""JARVIS background tasks."""
|
|
import asyncio
|
|
from typing import Optional
|
|
|
|
import rook
|
|
from dis_snek import Intents, Snake
|
|
from jarvis_core.db import connect
|
|
from jarvis_core.log import get_logger
|
|
|
|
from jarvis_tasks import const
|
|
from jarvis_tasks.config import TaskConfig
|
|
from jarvis_tasks.tasks import (
|
|
autokick,
|
|
ban,
|
|
lock,
|
|
lockdown,
|
|
reddit,
|
|
reminder,
|
|
twitter,
|
|
warning,
|
|
)
|
|
|
|
__version__ = const.__version__
|
|
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)
|
|
|
|
if config.rook_token:
|
|
rook.start(token=config.rook_token, labels={"env": "dev"})
|
|
|
|
# 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 = [
|
|
autokick.autokick,
|
|
ban.unban,
|
|
lock.unlock,
|
|
lockdown.lift,
|
|
reddit.reddit,
|
|
reminder.remind,
|
|
twitter.twitter,
|
|
warning.unwarn,
|
|
]
|
|
tasks = [loop.create_task(f(bot)) 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))
|