Add deadman switch cog
This commit is contained in:
@@ -0,0 +1,238 @@
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Optional
|
||||
|
||||
import discord
|
||||
from redbot.core import Config, commands
|
||||
|
||||
|
||||
def parse_duration(text: str) -> Optional[timedelta]:
|
||||
"""Parse a flexible duration string like ``6h``, ``2h30m``, ``90s`` or ``1d``."""
|
||||
text = text.strip().lower()
|
||||
if not text:
|
||||
return None
|
||||
|
||||
total = timedelta()
|
||||
num = ""
|
||||
unit = ""
|
||||
seen_unit = False
|
||||
|
||||
for char in text:
|
||||
if char.isdigit():
|
||||
if seen_unit:
|
||||
# A digit after a unit means a new component, e.g. "1d12h"
|
||||
delta = _unit_to_timedelta(int(num), unit)
|
||||
if delta is None:
|
||||
return None
|
||||
total += delta
|
||||
num = char
|
||||
unit = ""
|
||||
seen_unit = False
|
||||
else:
|
||||
num += char
|
||||
else:
|
||||
if not num:
|
||||
return None
|
||||
unit += char
|
||||
seen_unit = True
|
||||
|
||||
if unit:
|
||||
delta = _unit_to_timedelta(int(num), unit)
|
||||
if delta is None:
|
||||
return None
|
||||
total += delta
|
||||
elif num:
|
||||
# Trailing digits with no unit, e.g. "5m30" - invalid.
|
||||
return None
|
||||
|
||||
if total <= timedelta(0):
|
||||
return None
|
||||
|
||||
return total
|
||||
|
||||
|
||||
def _unit_to_timedelta(amount: int, unit: str) -> Optional[timedelta]:
|
||||
if not amount or not isinstance(amount, int):
|
||||
return None
|
||||
if unit and unit.isdigit():
|
||||
return None
|
||||
if unit in ("s", "sec", "secs", "second", "seconds"):
|
||||
return timedelta(seconds=amount)
|
||||
if unit in ("m", "min", "mins", "minute", "minutes"):
|
||||
return timedelta(minutes=amount)
|
||||
if unit in ("h", "hr", "hrs", "hour", "hours"):
|
||||
return timedelta(hours=amount)
|
||||
if unit in ("d", "day", "days"):
|
||||
return timedelta(days=amount)
|
||||
if unit in ("w", "week", "weeks"):
|
||||
return timedelta(weeks=amount)
|
||||
return None
|
||||
|
||||
|
||||
class Deadman(commands.Cog):
|
||||
"""Deadman switch - pings listed users if [p]deadman disable isn't run before expiration."""
|
||||
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
self.config = Config.get_conf(self, identifier=918273645)
|
||||
self.config.register_global(
|
||||
armed=False,
|
||||
expires_at=None,
|
||||
users=[],
|
||||
trigger_channel=None,
|
||||
)
|
||||
self._task = None
|
||||
self._triggered = False
|
||||
|
||||
async def red_delete_data_for_user(self, **kwargs):
|
||||
return
|
||||
|
||||
def cog_unload(self):
|
||||
if self._task:
|
||||
self._task.cancel()
|
||||
|
||||
async def cog_load(self):
|
||||
self._task = asyncio.create_task(self._monitor())
|
||||
|
||||
async def _monitor(self):
|
||||
"""Background task that checks every 30 seconds whether the switch has expired."""
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(30)
|
||||
await self._check_expiry()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _check_expiry(self):
|
||||
if not await self.config.armed():
|
||||
self._triggered = False
|
||||
return
|
||||
|
||||
expires_at = await self.config.expires_at()
|
||||
if expires_at is None:
|
||||
return
|
||||
|
||||
if datetime.utcnow() >= datetime.fromisoformat(expires_at):
|
||||
if not self._triggered:
|
||||
await self._trigger()
|
||||
# Keep armed/stale so repeated pings don't spam; disable clears state.
|
||||
|
||||
async def _trigger(self):
|
||||
self._triggered = True
|
||||
users = await self.config.users()
|
||||
channel_id = await self.config.trigger_channel()
|
||||
|
||||
victims: list[discord.User] = []
|
||||
for uid in users:
|
||||
user = self.bot.get_user(uid)
|
||||
if user is None:
|
||||
try:
|
||||
user = await self.bot.fetch_user(uid)
|
||||
except Exception:
|
||||
continue
|
||||
if user is not None:
|
||||
victims.append(user)
|
||||
|
||||
for user in victims:
|
||||
try:
|
||||
await user.send(
|
||||
"**GENERAL ALERT**\nThe deadman switch has been TRIGGERED. "
|
||||
"Please contact the operator immediately."
|
||||
)
|
||||
except discord.Forbidden:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if channel_id:
|
||||
channel = self.bot.get_channel(channel_id)
|
||||
if channel:
|
||||
mentions = " ".join(
|
||||
f"<@{uid}>" for uid in users if self.bot.get_user(uid)
|
||||
)
|
||||
try:
|
||||
await channel.send(
|
||||
f"**GENERAL ALERT**\nThe deadman switch has been TRIGGERED. "
|
||||
f"{mentions} Please contact the operator immediately."
|
||||
)
|
||||
except discord.Forbidden:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@commands.group()
|
||||
@commands.is_owner()
|
||||
async def deadman(self, ctx):
|
||||
"""Deadman switch control."""
|
||||
|
||||
@deadman.command(name="arm")
|
||||
async def deadman_arm(self, ctx, duration: str, users: commands.Greedy[discord.User], trigger_channel: Optional[discord.TextChannel] = None):
|
||||
"""Arm the deadman switch for a duration, pinging listed users if not disabled in time.
|
||||
|
||||
Duration format: a flexible string like ``6h``, ``2h30m``, ``90s``, ``1d``, or a combination.
|
||||
Users: mention the users to alert (e.g. ``@user1 @user2``).
|
||||
Optional: a channel where the alert should also be posted.
|
||||
"""
|
||||
if not users:
|
||||
await ctx.send("You must provide at least one user to alert.")
|
||||
return
|
||||
|
||||
duration_td = parse_duration(duration)
|
||||
if duration_td is None:
|
||||
await ctx.send(
|
||||
"Invalid duration. Examples: `6h`, `2h30m`, `90s`, `1d`, `1d12h`."
|
||||
)
|
||||
return
|
||||
|
||||
expires_at = datetime.utcnow() + duration_td
|
||||
await self.config.armed.set(True)
|
||||
await self.config.expires_at.set(expires_at.isoformat())
|
||||
await self.config.users.set([u.id for u in users])
|
||||
await self.config.trigger_channel.set(trigger_channel.id if trigger_channel else None)
|
||||
self._triggered = False
|
||||
|
||||
expires_str = expires_at.strftime("%Y-%m-%d %H:%M UTC")
|
||||
await ctx.send(
|
||||
f"Deadman switch ARMED.\n"
|
||||
f"Expires: **{expires_str}** (in {duration_td})\n"
|
||||
f"Users notified: {', '.join(u.mention for u in users)}\n"
|
||||
+ (f"Alert channel: {trigger_channel.mention}\n" if trigger_channel else "\n")
|
||||
+ "Run `[p]deadman disable` before that time to cancel."
|
||||
)
|
||||
|
||||
@deadman.command(name="disable")
|
||||
async def deadman_disable(self, ctx):
|
||||
"""Disable the currently armed deadman switch."""
|
||||
armed = await self.config.armed()
|
||||
if not armed:
|
||||
await ctx.send("The deadman switch is not currently armed.")
|
||||
return
|
||||
|
||||
await self.config.armed.set(False)
|
||||
await self.config.expires_at.set(None)
|
||||
await self.config.users.set([])
|
||||
await self.config.trigger_channel.set(None)
|
||||
self._triggered = False
|
||||
await ctx.send("Deadman switch disabled.")
|
||||
|
||||
@deadman.command(name="status")
|
||||
async def deadman_status(self, ctx):
|
||||
"""Show the current deadman switch status."""
|
||||
armed = await self.config.armed()
|
||||
if not armed:
|
||||
await ctx.send("The deadman switch is currently **disarmed**.")
|
||||
return
|
||||
|
||||
expires_at = await self.config.expires_at()
|
||||
users = await self.config.users()
|
||||
expires_dt = datetime.fromisoformat(expires_at)
|
||||
remaining = expires_dt - datetime.utcnow()
|
||||
mentions = ", ".join(f"<@{uid}>" for uid in users) if users else "None"
|
||||
await ctx.send(
|
||||
f"Deadman switch is **ARMED**.\n"
|
||||
f"Expires at: {expires_dt.strftime('%Y-%m-%d %H:%M UTC')}\n"
|
||||
f"Time remaining: {remaining}\n"
|
||||
f"Users to alert: {mentions}"
|
||||
)
|
||||
Reference in New Issue
Block a user