[Plugins/Jail] Save jail sentences to disk!

plugins/Jail/JailPlugin.py

@@ -1,10 +1,13 @@
 from argparse import ArgumentError
 import random
+import pickle
+import os
 
 import maya
 from requests.exceptions import HTTPError
 
 from plugins import PluginBase
+from core.Config import CACHE_DIR
 from plugins.Command import Command, CommandTypes
 from events.Message import Message, SysMessage, USER_TYPE
 from owoify import owoify
@@ -27,8 +30,12 @@ release_msgs = [
 
 
 class JailPlugin(PluginBase):
-    def __init__(self, *args, min_level='vip', **kwargs):
+    def __init__(self, *args, min_level='vip', persist=True, **kwargs):
         super().__init__(*args, **kwargs)
+        self.persist = persist
+        self._cache = os.path.join(CACHE_DIR, 'Jail', 'sentences')
+        os.makedirs(os.path.dirname(self._cache), exist_ok=True)
+
         self.jail_command = Command(name="jail", help='weewoo mowerwucker', required_level=min_level)
         self.jail_command.add_arg('username', type=CommandTypes.Username)
         self.jail_command.add_arg('--type', '-t', choices=['owo', 'uwu', 'uvu', 'bean'], default='uwu')
@@ -41,11 +48,29 @@ class JailPlugin(PluginBase):
         self.plugins['command'].register_help(self.unjail_command)
 
         self.sentences = {}
+        if self.persist and os.path.exists(self._cache):
+            with open(self._cache, 'rb') as pf:
+                saved = pickle.load(pf)
+            self.load_from(saved)
+
+    def on_change(self):
+        if self.persist:
+            with open(self._cache, 'wb') as pf:
+                pickle.dump(self.sentences, pf)
+
+    def load_from(self, sentences):
+        current_time = maya.now()
+        for key, value in sentences.items():
+            end_date, type, username = value
+            if current_time < end_date:
+                self.sentences[key] = value
 
     def tick(self, dt):
+        current_time = maya.now()
         for key, (end_date, type, username) in list(self.sentences.items())[:]:
-            if maya.now() > end_date:
+            if current_time > end_date:
                 del self.sentences[key]
+                self.on_change()
                 msg = random.choice(release_msgs).format(user=username)
                 freedom = SysMessage(self._name, msg)
                 self.broadcast(freedom)
@@ -56,13 +81,7 @@ class JailPlugin(PluginBase):
 
     def on_bus_event(self, event):
         if isinstance(event, Message):
-            if sentence := self.sentences.get(event.user_id):
-                end_date, type, username = sentence
-                if type in ['owo', 'uwu', 'uvu']:
-                    event.text = owoify(event.text, type)
-                elif type == 'bean':
-                    event.text = ' '.join(['bean'] * len(event.text.split()))
-            elif self.jail_command.invoked(event):
+            if self.jail_command.invoked(event):
                 try:
                     args = self.jail_command.parse(event.text)
                     end_date = maya.when(args['length'], prefer_dates_from='future')
@@ -80,8 +99,13 @@ class JailPlugin(PluginBase):
                     self.logger.debug(e);
                     return None
 
-
-                self.sentences[user_id] = (end_date, args['type'], args['username'])
+                if self.sentences.get(user_id) is not None:
+                    msg = SysMessage(self._name, "Jail fail - user is already in jail!", replies_to=event)
+                    self.chats[event.via].send(msg)
+                    return None
+                else:
+                    self.sentences[user_id] = (end_date, args['type'], args['username'])
+                    self.on_change()
 
                 msg = random.choice(admition_msgs).format(user=args['username'])
                 weewoo = SysMessage(self._name, msg)
@@ -95,6 +119,7 @@ class JailPlugin(PluginBase):
                         raise ValueError()
                     user_id = str(deets['user']['id'])
                     del self.sentences[user_id]
+                    self.on_change()
                 except ArgumentError as e:
                     msg = SysMessage(self._name, str(e), replies_to=event)
                     self.chats[event.via].send(msg)
@@ -108,5 +133,11 @@ class JailPlugin(PluginBase):
                 weewoo = SysMessage(self._name, msg)
                 self.broadcast(weewoo)
                 self.send_to_bus(weewoo)
+            elif sentence := self.sentences.get(event.user_id):
+                end_date, type, username = sentence
+                if type in ['owo', 'uwu', 'uvu']:
+                    event.text = owoify(event.text, type)
+                elif type == 'bean':
+                    event.text = ' '.join(['bean'] * len(event.text.split()))
 
         return event