From: Thorsten <mail@aero2k.de>
Date: Sat, 7 Sep 2024 12:14:10 +0000 (+0200)
Subject: add poll plugin
X-Git-Url: https://git.aero2k.de/?a=commitdiff_plain;h=369cf5bc2b9f08d94ddb34a13bbacbfcfd15d8ef;p=urlbot-v3.git

add poll plugin
---

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]