mirror of
https://github.com/nexy7574/LCC-bot.git
synced 2024-09-19 10:03:40 +01:00
Persistence!
This commit is contained in:
parent
49b425163a
commit
84fdd21a9d
GPG key ID:
0FA334385D0B689F
1 changed files with 131 additions and 39 deletions
|
|
@ -1,5 +1,7 @@
|
|||
import asyncio
|
||||
import pathlib
|
||||
import re
|
||||
import typing
|
||||
|
||||
import aiosqlite
|
||||
|
||||
|
|
@ -7,12 +9,101 @@ import discord
|
|||
from discord.ext import commands
|
||||
|
||||
|
||||
class McDataBase:
|
||||
def __init__(self):
|
||||
self.db = pathlib.Path.home() / ".cache" / "McDataBase.db"
|
||||
self._conn: typing.Optional[aiosqlite.Connection] = None
|
||||
|
||||
async def init_db(self):
|
||||
if self._conn:
|
||||
conn = self._conn
|
||||
await conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS breaks (
|
||||
user_id INTEGER PRIMARY KEY,
|
||||
since FLOAT NOT NULL,
|
||||
author INTEGER NOT NULL
|
||||
);
|
||||
"""
|
||||
)
|
||||
await conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS cooldowns (
|
||||
user_id INTEGER PRIMARY KEY,
|
||||
expires FLOAT NOT NULL
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
async def get_break(self, user_id: int) -> typing.Optional[float]:
|
||||
async with self._conn.execute(
|
||||
"""
|
||||
SELECT since FROM breaks WHERE user_id = ?;
|
||||
""",
|
||||
(user_id,)
|
||||
) as cursor:
|
||||
return await cursor.fetchone()
|
||||
|
||||
async def set_break(self, user_id: int, since: float):
|
||||
await self._conn.execute(
|
||||
"""
|
||||
INSERT INTO breaks (user_id, since) VALUES (?, ?)
|
||||
ON CONFLICT(user_id) DO UPDATE SET since = excluded.since
|
||||
""",
|
||||
(user_id, since)
|
||||
)
|
||||
|
||||
async def remove_break(self, user_id: int):
|
||||
now = discord.utils.utcnow().timestamp()
|
||||
await self._conn.execute(
|
||||
"""
|
||||
DELETE FROM breaks WHERE user_id = ?;
|
||||
""",
|
||||
(user_id,)
|
||||
)
|
||||
await self.set_cooldown(user_id, now)
|
||||
|
||||
async def get_cooldown(self, user_id: int) -> typing.Optional[float]:
|
||||
async with self._conn.execute(
|
||||
"""
|
||||
SELECT expires FROM cooldowns WHERE user_id = ?;
|
||||
""",
|
||||
(user_id,)
|
||||
) as cursor:
|
||||
return await cursor.fetchone()
|
||||
|
||||
async def set_cooldown(self, user_id: int, expires: float):
|
||||
await self._conn.execute(
|
||||
"""
|
||||
INSERT INTO cooldowns (user_id, expires) VALUES (?, ?)
|
||||
ON CONFLICT(user_id) DO UPDATE SET expires = excluded.expires;
|
||||
""",
|
||||
(user_id, expires)
|
||||
)
|
||||
|
||||
async def remove_cooldown(self, user_id: int):
|
||||
await self._conn.execute(
|
||||
"""
|
||||
DELETE FROM cooldowns WHERE user_id = ?;
|
||||
""",
|
||||
(user_id,)
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> "McDataBase":
|
||||
self._conn = await aiosqlite.connect(self.db)
|
||||
await self.init_db()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
await self._conn.commit()
|
||||
await self._conn.close()
|
||||
self._conn = None
|
||||
|
||||
|
||||
class McDonaldsCog(commands.Cog):
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
self.targets: dict[discord.Member, float] = {}
|
||||
self.lock = asyncio.Lock()
|
||||
self.cooldown: dict[discord.Member, float] = {}
|
||||
|
||||
@commands.Cog.listener()
|
||||
async def on_message(self, message: discord.Message):
|
||||
|
|
@ -22,29 +113,29 @@ class McDonaldsCog(commands.Cog):
|
|||
return
|
||||
|
||||
async with self.lock:
|
||||
NIGHTMARE_REGEX = re.compile(r"\|\| \|\|([^|]+)#0\|\| \|\|: bypassed.*")
|
||||
NIGHTMARE_REGEX = re.compile(r"[|\s]*(?P<username>[^|]+)(#0)?[|\s]*:( bypassed.*)?")
|
||||
if m := NIGHTMARE_REGEX.match(message.content):
|
||||
username = m.group(1)
|
||||
member = discord.utils.get(message.guild.members, name=username)
|
||||
if member:
|
||||
author = member
|
||||
|
||||
if author in self.targets:
|
||||
if message.content.upper() != "MCDONALDS!":
|
||||
await message.delete()
|
||||
if (message.created_at.timestamp() - self.targets[author]) > 10:
|
||||
await message.channel.send(
|
||||
f"{message.author.mention} Please say `MCDONALDS!` to end commercial.",
|
||||
delete_after=30
|
||||
async with McDataBase() as db:
|
||||
if (last_info := await db.get_break(author.id)) is not None:
|
||||
if message.content.upper() != "MCDONALDS!":
|
||||
await message.delete()
|
||||
if (message.created_at.timestamp() - last_info) > 10:
|
||||
await message.channel.send(
|
||||
f"{message.author.mention} Please say `MCDONALDS!` to end commercial.",
|
||||
delete_after=30
|
||||
)
|
||||
await db.set_break(author.id, message.created_at.timestamp())
|
||||
else:
|
||||
await db.remove_break(author.id)
|
||||
await message.reply(
|
||||
"Thank you. You may now resume your activity.",
|
||||
delete_after=120
|
||||
)
|
||||
self.targets[author] = message.created_at.timestamp()
|
||||
else:
|
||||
self.targets.pop(author, None)
|
||||
self.cooldown[author] = message.created_at.timestamp()
|
||||
await message.reply(
|
||||
"Thank you. You may now resume your activity.",
|
||||
delete_after=120
|
||||
)
|
||||
|
||||
@commands.user_command(name="Commercial Break")
|
||||
async def commercial_break(self, ctx: discord.ApplicationContext, member: discord.Member):
|
||||
|
|
@ -56,29 +147,30 @@ class McDonaldsCog(commands.Cog):
|
|||
if member.bot or member == ctx.user:
|
||||
return await ctx.respond("No.", ephemeral=True)
|
||||
|
||||
if member in self.targets.keys():
|
||||
await ctx.respond(f"{member.mention} is already in a commercial break.")
|
||||
return
|
||||
elif member in self.cooldown.keys():
|
||||
expires = self.cooldown[member] + 300
|
||||
if expires > discord.utils.utcnow().timestamp():
|
||||
await ctx.respond(
|
||||
f"{member.mention} is not due another ad break yet. Their next commercial break will start "
|
||||
f"<t:{int(expires)}:R> at the earliest."
|
||||
)
|
||||
async with McDataBase() as db:
|
||||
if await db.get_break(member.id) is not None:
|
||||
await ctx.respond(f"{member.mention} is already in a commercial break.")
|
||||
return
|
||||
else:
|
||||
self.cooldown.pop(member, None)
|
||||
elif (cooldown := await db.get_cooldown(member.id)) is not None:
|
||||
expires = cooldown + 300
|
||||
if expires > discord.utils.utcnow().timestamp():
|
||||
await ctx.respond(
|
||||
f"{member.mention} is not due another ad break yet. Their next commercial break will start "
|
||||
f"<t:{int(expires)}:R> at the earliest."
|
||||
)
|
||||
return
|
||||
else:
|
||||
await db.remove_cooldown(member.id)
|
||||
|
||||
self.targets[member] = discord.utils.utcnow().timestamp()
|
||||
await ctx.send(
|
||||
f"{member.mention} Commercial break! Please say `MCDONALDS!` to end commercial.\n"
|
||||
f"*This commercial break is sponsored by {ctx.user.mention}.*",
|
||||
delete_after=300,
|
||||
allowed_mentions=discord.AllowedMentions(users=True, roles=False, everyone=False)
|
||||
)
|
||||
await ctx.respond("Commercial break started.", ephemeral=True)
|
||||
await ctx.delete(delay=120)
|
||||
await db.set_break(member.id, discord.utils.utcnow().timestamp())
|
||||
await ctx.send(
|
||||
f"{member.mention} Commercial break! Please say `MCDONALDS!` to end commercial.\n"
|
||||
f"*This commercial break is sponsored by {ctx.user.mention}.*",
|
||||
delete_after=300,
|
||||
allowed_mentions=discord.AllowedMentions(users=True, roles=False, everyone=False)
|
||||
)
|
||||
await ctx.respond("Commercial break started.", ephemeral=True)
|
||||
await ctx.delete(delay=120)
|
||||
|
||||
|
||||
def setup(bot):
|
||||
|
|
|
|||
Loading…
Reference in a new issue