From 369cf5bc2b9f08d94ddb34a13bbacbfcfd15d8ef Mon Sep 17 00:00:00 2001 From: Thorsten Date: Sat, 7 Sep 2024 14:14:10 +0200 Subject: [PATCH] add poll plugin --- pyproject.toml | 4 +- src/distbot/bot/worker.py | 2 +- src/distbot/minijobber/run.py | 4 +- src/distbot/plugins/fun.py | 1 - src/distbot/plugins/votepoll.py | 193 ++++++++++++++++++++++++++++++++ 5 files changed, 199 insertions(+), 5 deletions(-) create mode 100644 src/distbot/plugins/votepoll.py diff --git a/pyproject.toml b/pyproject.toml index 78f73c1..9baac8c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,8 +22,8 @@ dependencies = [ ] [project.optional-dependencies] -chatbot = [] # "pyopenssl" should be dropped and if required, replaced with pyca/cryptography. -worker = ["lxml"] +chatbot = [] +worker = ["lxml", "plyvel"] [project.scripts] urlbotd-chat = "distbot.bot.bot:run" diff --git a/src/distbot/bot/worker.py b/src/distbot/bot/worker.py index ec69805..4e348d1 100644 --- a/src/distbot/bot/worker.py +++ b/src/distbot/bot/worker.py @@ -121,7 +121,7 @@ class Worker(threading.Thread): self.channel.basic_consume(queue=self.queue, on_message_callback=self.callback) self.channel.start_consuming() - def parse_body(self, msg): + def parse_body(self, msg: dict) -> Action | None: raise NotImplementedError() def register_plugin(self): diff --git a/src/distbot/minijobber/run.py b/src/distbot/minijobber/run.py index 5044a20..d8e1ff9 100644 --- a/src/distbot/minijobber/run.py +++ b/src/distbot/minijobber/run.py @@ -12,7 +12,8 @@ from distbot.plugins import ( basic, fun, lookup, url, feeds, muc, translation, searx, queue_management, plugin_help, morse, meta, extended, bugtracker, bots, bofh, didyouknow, - youtube + votepoll, + youtube, ) logger = logging.getLogger(__name__) @@ -40,6 +41,7 @@ PLUGIN_MODULES = { didyouknow: didyouknow.ALL, youtube: youtube.ALL, # debug: debug.ALL + votepoll: votepoll.ALL, } job_workers: list[Worker] = [] diff --git a/src/distbot/plugins/fun.py b/src/distbot/plugins/fun.py index ace0868..1034a76 100644 --- a/src/distbot/plugins/fun.py +++ b/src/distbot/plugins/fun.py @@ -1,6 +1,5 @@ # -*- coding: utf-8 -*- import logging - import random import time diff --git a/src/distbot/plugins/votepoll.py b/src/distbot/plugins/votepoll.py new file mode 100644 index 0000000..06852d2 --- /dev/null +++ b/src/distbot/plugins/votepoll.py @@ -0,0 +1,193 @@ +import hashlib +import json +import time +from dataclasses import dataclass +from datetime import datetime + +import plyvel + +from distbot.bot.worker import Worker +from distbot.common.action import Action +from distbot.common.config import conf_get +from distbot.common.message import get_words, get_nick_from_message + + +@dataclass +class Poll: + option_a: str + option_b: str + votes: dict[str, str] + ending_at: float + + def add_vote(self, user, vote): + # TODO return copy of itself or use mutable object? + # TODO maybe forbid changing vote + if vote in self.options: + self.votes[user] = vote + return self + + @property + def options(self): + return {self.option_a, self.option_b} + + @property + def key(self): + return self.generate_key(self.option_a, self.option_b) + + def votes_a(self) -> int: + return len([vote for vote in self.votes.values() if vote == self.option_a]) + + def votes_b(self) -> int: + return len([vote for vote in self.votes.values() if vote == self.option_b]) + + @staticmethod + def generate_key(option_a, option_b): + return hashlib.sha256("".join(sorted((option_a, option_b))).encode()).hexdigest() + + def due(self, timestamp=None) -> bool: + if not timestamp: + timestamp = time.time() + return timestamp > self.ending_at + + @staticmethod + def from_json(data: dict): + try: + return Poll( + option_a=data['option_a'], + option_b=data['option_b'], + votes=data['votes'], + ending_at=float(data['ending_at']) + ) + except (KeyError, TypeError, ValueError) as e: + raise ValueError(f"Invalid poll data: {e}") + + def to_json(self): + return json.dumps(self, default=lambda o: o.__dict__) + + def status_report(self): + # Calculate winner and format result message + winner = "It's a tie!" + if self.votes_a() > self.votes_b(): + winner = f"**Winner:** {self.option_a}" + elif self.votes_a() < self.votes_b(): + winner = f"**Winner:** {self.option_b}" + + status_message = (f"**Poll status:**\n" + f" * {self.option_a}: {self.votes_a()}\n" + f" * {self.option_b}: {self.votes_b()}" + ) + (f"\n{winner}" if self.due() else "") + return status_message + + +class VotePoll(Worker): + """ + leveldb supported voting plugin + leveldb contents: + - hashed key for each poll + - "active" key with the hashed key of the active poll + """ + binding_keys = ["nick.poll.*.vs.*", "nick.vote.*", "nick.pollstatus"] + description = "Vote on polls (A vs B)" + + vote_duration = 60 * 1 + + KEY_ACTIVE = "active".encode() + + def __init__(self, actionqueue, queue="work"): + super().__init__(actionqueue, queue) + self.db = plyvel.DB("votepoll.db", create_if_missing=True) + + def parse_body(self, msg: dict) -> Action | None: + sender = get_nick_from_message(msg) + words = get_words(msg) + now: float = time.time() + + # check active polls for timeout, disabling any left-overs + active_poll = self.get_active_poll() + if active_poll and active_poll.due() and words[0] != "pollstatus": + # something went wrong - only pollstatus should ever end the poll. + self.close_poll(active_poll.key) + + if words[0] == "vote": + # check if words contain any option + map_ab = { + "A": active_poll.option_a, + "B": active_poll.option_b, + } + vote = map_ab.get(words[1].upper(), words[1]) + + if vote in active_poll.options: + self.persist(active_poll.add_vote(sender, vote)) + return Action(msg="Vote added.") + else: + return Action(msg=f"not a valid option: {vote} (valid: {active_poll.options})") + elif words[0] == "poll": + # parse options + option_a = words[1] + option_b = words[3] + # setup new poll + poll = Poll(option_a, option_b, {}, now + self.vote_duration) + # check prior results + if active_poll: + return Action(msg="There is already an active poll.") + else: + prior_poll = self.get_poll(poll.key) + if prior_poll and prior_poll.votes: + date = datetime.fromtimestamp(prior_poll.ending_at) + return Action(msg=f"This was already decided at {date}: {prior_poll.status_report()}") + + # enable poll + self.start_poll(poll) + # setup timeout to disable poll and present results + bot_nick = conf_get("bot_nickname") + poll_message = f"**New Vote:** {sender} started a vote! Vote for A: {option_a} or B: {option_b} (reply with '{bot_nick}: vote foo' within {self.vote_duration}s)" + return Action( + msg=poll_message, + event=Action(time=now + self.vote_duration, command="nick.pollstatus", + msg=f'{sender}: Poll is over!', mutex=f'poll') + ) + elif words[0] == "pollstatus": + return self.check_poll() + + def get_active_poll(self) -> Poll: + active_key = self.db.get(self.KEY_ACTIVE) + if active_key: + return self.get_poll(active_key) + + def get_poll(self, key: str): + poll_bin = self.db.get(key.encode()) + poll = json.loads(poll_bin.decode("utf-8")) + return Poll.from_json(poll) + + def persist(self, poll: Poll): + self.db.put(poll.key.encode(), poll.to_json().encode()) + + def start_poll(self, poll: Poll): + key = poll.key.encode() + self.db.put(key, poll.to_json().encode()) + self.db.put(self.KEY_ACTIVE, key) + + def close_poll(self, key: str): + active_key = self.db.get(self.KEY_ACTIVE) + if active_key == key.encode(): + self.db.delete(self.KEY_ACTIVE) + + def check_poll(self): + poll = self.get_active_poll() + if not poll: + return Action(msg="no active poll") + if poll.due(): + return self.end_poll() + else: + return Action(msg=poll.status_report()) + + def end_poll(self) -> Action | None: + poll = self.get_active_poll() + if not poll: + return + self.close_poll(poll.key) + report = poll.status_report() + return Action(msg=report) + + +ALL = [VotePoll] -- 2.39.2