-
Notifications
You must be signed in to change notification settings - Fork 0
/
bot.py
115 lines (87 loc) · 2.95 KB
/
bot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
from __future__ import annotations
import aiob2
import datetime
import logging
from collections import Counter
from typing import TYPE_CHECKING, TypedDict
import aiohttp
import asyncpg
import discord
from discord.ext import commands
from utils.context import Context
from utils.sonarr import Client as SonarrClient
from cogs.youtube import DownloadControls
if TYPE_CHECKING:
from cogs.reminder import Reminder
class PostgresConfig(TypedDict):
host: str
user: str
password: str
database: str
class BackblazeConfig(TypedDict):
key_id: str
key: str
bucket_id: str
class SonarrConfig(TypedDict):
host: str
api_key: str
class WebhookConfig(TypedDict):
wh_id: int
wh_token: str
class Config(TypedDict):
token: str
owner_ids: list[int]
database: PostgresConfig
backblaze: BackblazeConfig
sonarr: SonarrConfig
error: WebhookConfig
class RoboDan(commands.Bot):
pool: asyncpg.Pool
session: aiohttp.ClientSession
bucket: aiob2.Client
sonarr: SonarrClient
command_stats: Counter[str]
socket_stats: Counter[str]
launch_time: datetime.datetime
def __init__(self, config: Config):
super().__init__(
command_prefix=commands.when_mentioned_or('-'),
intents=discord.Intents.all(),
owner_ids=config['owner_ids'],
chunk_guilds_at_startup=True
)
self.config = config
self.add_check(self.ctx_check)
self.add_view(DownloadControls())
self.tree.interaction_check = self.interaction_check
self.global_log = logging.getLogger('robodan')
@discord.utils.cached_property
def error_webhook(self):
hook = discord.Webhook.partial(
id=self.config['error']['wh_id'],
token=self.config['error']['wh_token'],
session=self.session
)
return hook
@property
def reminder(self) -> Reminder | None:
return self.get_cog('Reminder') # type: ignore
async def ctx_check(self, ctx: Context) -> bool:
return ctx.author.id == 723943620054614047
async def interaction_check(self, interaction: discord.Interaction) -> bool:
return interaction.user.id == 723943620054614047
async def get_context(self, message, *, cls=Context):
return await super().get_context(message, cls=cls)
async def startup_message(self):
await self.wait_until_ready()
self.global_log.info('Bot is ready with a populated cache')
async def setup_hook(self) -> None:
self.launch_time = datetime.datetime.utcnow()
async def start(self, *, reconnect: bool = True) -> None:
await super().start(self.config['token'], reconnect=reconnect)
async def close(self) -> None:
await super().close()
if hasattr(self, 'pool') and self.pool is not None:
await self.pool.close()
if hasattr(self, 'session') and self.session is not None:
await self.session.close()