125 lines
4.2 KiB
Python
125 lines
4.2 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
|
|
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._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):
|
|
targets = await discover_targets(self, self.cfg.friend_id, 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)
|
|
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,
|
|
)
|
|
|
|
for target in targets:
|
|
if target.target_type == "group_dm":
|
|
channel = self.get_channel(target.target_id)
|
|
if channel is None:
|
|
log.warning("Group DM %s not found in cache, skipping", target.target_id)
|
|
continue
|
|
self.known_channel_targets[channel.id] = target
|
|
await scan_channel_backlog(self.storage, channel, target, self.cfg.friend_id, cutoff, react_fn)
|
|
else:
|
|
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, cutoff, react_fn)
|
|
|
|
log.info("Backlog scan complete. Now watching live.")
|
|
|
|
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:
|
|
return
|
|
if message.author.id != self.cfg.friend_id:
|
|
return
|
|
await react_to_message(
|
|
self.storage, message, target, self.cfg.friend_id, self.cfg.emoji,
|
|
source="live", display=self.display,
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
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()
|