Minz
52b0fda829
on_message only matched channels enumerated during the backlog scan, so a brand-new channel (or any type not scanned, like threads) was silently ignored until the next restart. Now falls back to a guild-level lookup and registers the channel for live watching going forward the first time a message arrives there.
175 lines
6.9 KiB
Python
175 lines
6.9 KiB
Python
import argparse
|
|
import asyncio
|
|
import datetime
|
|
import logging
|
|
|
|
import discord
|
|
from rich.console import Console
|
|
from rich.logging import RichHandler
|
|
|
|
from bot.backlog import scan_channel_backlog
|
|
from bot.channels import get_scannable_channels
|
|
from bot.config import load_config
|
|
from bot.discovery import Target, discover_targets, resolve_friend
|
|
from bot.display import StatsDisplay
|
|
from bot.reactor import react_to_message
|
|
from bot.storage import Storage
|
|
from bot.undo import undo_all_reactions
|
|
|
|
log = logging.getLogger("dc_reactor")
|
|
|
|
|
|
class ReactorClient(discord.Client):
|
|
def __init__(self, cfg, storage: Storage, display: StatsDisplay, undo: bool = False, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.cfg = cfg
|
|
self.storage = storage
|
|
self.display = display
|
|
self.undo = undo
|
|
self.known_channel_targets: dict[int, Target] = {}
|
|
self.known_guild_targets: dict[int, Target] = {}
|
|
self._started = False
|
|
|
|
async def on_ready(self):
|
|
log.info("Logged in as %s (%s)", self.user, self.user.id)
|
|
if self._started:
|
|
return
|
|
self._started = True
|
|
if self.undo:
|
|
self.loop.create_task(self._run_undo())
|
|
else:
|
|
self.loop.create_task(self._run_backlog_then_watch())
|
|
|
|
async def _run_undo(self):
|
|
await undo_all_reactions(self, self.storage, display=self.display)
|
|
await self.close()
|
|
|
|
async def _run_backlog_then_watch(self):
|
|
self.display.set_status(f"Resolving friend id {self.cfg.friend_id}...")
|
|
try:
|
|
friend = await resolve_friend(self, self.cfg.friend_id)
|
|
except discord.NotFound:
|
|
log.error(
|
|
"friend_id %s does not resolve to any Discord user. Check config.json and fix it.",
|
|
self.cfg.friend_id,
|
|
)
|
|
self.display.set_status("Failed: friend_id does not resolve to a Discord user.")
|
|
return
|
|
|
|
log.info("Targeting friend: %s (%s)", friend, friend.id)
|
|
self.display.set_friend(str(friend))
|
|
self.display.set_status(f"Discovering mutual servers/groups with {friend}...")
|
|
await self.display.refresh()
|
|
|
|
targets = await discover_targets(self, friend, self.cfg.manual_guild_ids)
|
|
log.info("Discovered %d mutual target(s): %s", len(targets), [t.name for t in targets])
|
|
|
|
for t in targets:
|
|
await asyncio.to_thread(self.storage.upsert_target, t.target_id, t.target_type, t.name)
|
|
pruned = await asyncio.to_thread(self.storage.prune_targets, [t.target_id for t in targets])
|
|
if pruned:
|
|
log.info("Pruned %d stale target(s) no longer mutual/discovered", pruned)
|
|
self.display.set_status(f"Starting backlog scan across {len(targets)} target(s)...")
|
|
await self.display.refresh()
|
|
|
|
cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=self.cfg.days_back)
|
|
|
|
async def react_fn(message: discord.Message, target: Target) -> bool:
|
|
return await react_to_message(
|
|
self.storage, message, target, self.cfg.friend_id, self.cfg.emoji,
|
|
source="backlog", display=self.display, count_found=False,
|
|
)
|
|
|
|
for target in targets:
|
|
self.display.set_active_target(target.target_id)
|
|
if target.target_type in ("group_dm", "dm"):
|
|
channel = self.get_channel(target.target_id)
|
|
if channel is None:
|
|
log.warning("Channel %s (%s) not found in cache, skipping", target.target_id, target.name)
|
|
continue
|
|
self.known_channel_targets[channel.id] = target
|
|
await scan_channel_backlog(
|
|
self.storage, channel, target, self.cfg.friend_id, self.cfg.emoji,
|
|
cutoff, react_fn, display=self.display,
|
|
)
|
|
else:
|
|
self.known_guild_targets[target.target_id] = target
|
|
for channel in get_scannable_channels(self, target.target_id):
|
|
self.known_channel_targets[channel.id] = target
|
|
await scan_channel_backlog(
|
|
self.storage, channel, target, self.cfg.friend_id, self.cfg.emoji,
|
|
cutoff, react_fn, display=self.display,
|
|
)
|
|
|
|
self.display.set_active_target(None)
|
|
log.info("Backlog scan complete. Now watching live.")
|
|
self.display.set_status("Backlog scan complete — watching live for new messages...")
|
|
|
|
async def on_message(self, message: discord.Message):
|
|
if self.undo:
|
|
return
|
|
target = self.known_channel_targets.get(message.channel.id)
|
|
if target is None:
|
|
# Channel didn't exist (or wasn't enumerated, e.g. a thread) at
|
|
# backlog-scan time. If it belongs to an already-known mutual
|
|
# guild, start watching it now instead of requiring a restart —
|
|
# only its live messages going forward are covered, no backfill.
|
|
guild = getattr(message, "guild", None)
|
|
if guild is not None:
|
|
target = self.known_guild_targets.get(guild.id)
|
|
if target is not None:
|
|
self.known_channel_targets[message.channel.id] = target
|
|
log.info(
|
|
"New channel #%s in %s wasn't seen during backlog scan; watching it live from now on",
|
|
getattr(message.channel, "name", message.channel.id), target.name,
|
|
)
|
|
if target is None:
|
|
return
|
|
if message.author.id != self.cfg.friend_id:
|
|
return
|
|
chan_label = getattr(message.channel, "name", None) or target.name
|
|
self.display.set_active_target(target.target_id)
|
|
self.display.set_status(f"New message from friend in {chan_label}...")
|
|
await react_to_message(
|
|
self.storage, message, target, self.cfg.friend_id, self.cfg.emoji,
|
|
source="live", display=self.display,
|
|
)
|
|
self.display.set_active_target(None)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="React to a friend's messages across mutual servers/groups.")
|
|
parser.add_argument(
|
|
"--undo",
|
|
action="store_true",
|
|
help="Remove every reaction previously added by this script (per the local database log) and exit.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
cfg = load_config()
|
|
|
|
console = Console()
|
|
logging.basicConfig(
|
|
level=cfg.log_level,
|
|
format="%(message)s",
|
|
handlers=[RichHandler(console=console, show_path=False, rich_tracebacks=True)],
|
|
)
|
|
|
|
storage = Storage(cfg.db_path)
|
|
display = StatsDisplay(console, storage)
|
|
display.set_status("Connecting to Discord...")
|
|
|
|
with display:
|
|
client = ReactorClient(cfg, storage, display, undo=args.undo)
|
|
try:
|
|
client.run(cfg.token, log_handler=None)
|
|
finally:
|
|
storage.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|