Velasco 2.2

- Added optional configuration restrictions to admin-only

chatlog.py

@@ -11,7 +11,7 @@ def parse_line(l):
         return s[1]
 
 class Chatlog(object):
-    def __init__(self, ident, chattype, title, text=None, freq=None, answer=0.5):
+    def __init__(self, ident, chattype, title, text=None, freq=None, answer=0.5, restricted=False):
         self.id = str(ident)
         self.type = chattype
         self.title = title
@@ -28,6 +28,7 @@ class Chatlog(object):
             self.count = 0
         self.replyables = []
         self.answer = answer
+        self.restricted = restricted
         self.gen = Markov(text)
 
     def set_title(self, title):
@@ -71,6 +72,7 @@ class Chatlog(object):
             return False
         return rand <= self.answer
 
+<<<<<<< Updated upstream
     def add_replyable(self, msg_id):
         self.replyables.append(msg_id)
 
@@ -82,14 +84,22 @@ class Chatlog(object):
 
     def get_replyable(self):
         random.choice(self.replyables)
+=======
+    def toggle_restrict(self):
+        self.restricted = (not self.restricted)
+
+    def is_restricted(self):
+        return self.restricted
+>>>>>>> Stashed changes
 
     def to_txt(self):
-        lines = ["DICT=v2"]
+        lines = ["DICT=v3"]
         lines.append("CHAT_ID=" + self.id)
         lines.append("CHAT_TYPE=" + self.type)
         lines.append("CHAT_NAME=" + self.title)
         lines.append("MESSAGE_FREQ=" + str(self.freq))
         lines.append("ANSWER_FREQ=" + str(self.answer))
+        lines.append("RESTRICTED=" + str(self.restricted))
         lines.append("WORD_COUNT=" + str(self.count))
         lines.append("WORD_DICT=")
         txt = '\n'.join(lines)
@@ -99,7 +109,15 @@ class Chatlog(object):
         lines = text.splitlines()
         #print("Line 4=" + lines[4])
         print("Line 0=" + parse_line(lines[0]))
-        if(parse_line(lines[0]) == "v2"):
+        if(parse_line(lines[0]) == "v3"):
+            new_log = Chatlog(parse_line(lines[1]), parse_line(lines[2]), parse_line(lines[3]), None, int(parse_line(lines[4])), float(parse_line(lines[5])), (parse_line(lines[6]) == 'True'))
+            new_log.count = int(parse_line(lines[7]))
+            cache = '\n'.join(lines[9:])
+            new_log.gen = Markov.from_json(cache)
+            if new_log.count < 0:
+                new_log.count = new_log.gen.new_count()
+            return new_log
+        elif(parse_line(lines[0]) == "v2"):
             new_log = Chatlog(parse_line(lines[1]), parse_line(lines[2]), parse_line(lines[3]), None, int(parse_line(lines[4])), float(parse_line(lines[5])))
             new_log.count = int(parse_line(lines[6]))
             cache = '\n'.join(lines[8:])

velasco.py

@@ -76,7 +76,8 @@ def help(bot, update):
 /count - I tell you how many messages from this chat I remember.
 /freq - Change the frequency of my messages. (Maximum of 100000)
 /speak - Forces me to speak.
-/answer - Change the probability to answer to a reply. (Decimal between 0 and 1)
+/answer - Change the probability to answer to a reply. (Decimal between 0 and 1).
+/restrict - Toggle restriction of configuration commands to admins only.
     """)
 
 def about(bot, update):
@@ -171,6 +172,12 @@ def speak(bot, update):
         chatlog = Chatlog(chat.id, chat.type, title)
     else:
         chatlog = chatlogs[ident]
+
+    if chatlogs[ident].is_restricted():
+        user = update.message.chat.get_member(update.message.from_user.id)
+        if not user_is_admin(user):
+            return
+
     text = update.message.text.split()
     if len(text) > 1:
         chatlog.add_msg(' '.join(text[1:]))
@@ -213,6 +220,12 @@ def get_count(bot, update):
     reply += " messages."
     update.message.reply_text(reply)
 
+def user_is_admin(member):
+    global ADMIN_ID
+    print("user {} requesting a restricted action".format(str(member.user.id)))
+    print("Creator ID is {}".format(str(ADMIN_ID)))
+    return (member.status == 'creator') or (member.status == 'administrator') or (member.user.id == ADMIN_ID)
+
 def set_freq(bot, update):
     ident = str(update.message.chat.id)
     if not ident in chatlogs:
@@ -220,9 +233,16 @@ def set_freq(bot, update):
         title = get_chatname(chat)
         chatlog = Chatlog(chat.id, chat.type, title)
         chatlogs[chatlog.id] = chatlog
+
     if not len(update.message.text.split()) > 1:
         reply = "Current frequency is " + str(chatlogs[ident].freq)
     else:
+        if chatlogs[ident].is_restricted():
+            user = update.message.chat.get_member(update.message.from_user.id)
+            if not user_is_admin(user):
+                reply = "You do not have permissions to do that."
+                update.message.reply_text(reply)
+                return
         try:
             value = update.message.text.split()[1]
             value = int(value)
@@ -240,9 +260,16 @@ def set_answer_freq(bot, update):
         title = get_chatname(chat)
         chatlog = Chatlog(chat.id, chat.type, title)
         chatlogs[chatlog.id] = chatlog
+
     if not len(update.message.text.split()) > 1:
         reply = "Current answer probability is " + str(chatlogs[ident].answer)
     else:
+        if chatlogs[ident].is_restricted():
+            user = chat.get_member(update.message.from_user.id)
+            if not user_is_admin(user):
+                reply = "You do not have permissions to do that."
+                update.message.reply_text(reply)
+                return
         try:
             value = update.message.text.split()[1]
             value = float(value)
@@ -253,6 +280,28 @@ def set_answer_freq(bot, update):
             reply = "Format was confusing; answer probability not changed from " + str(chatlogs[ident].answer)
     update.message.reply_text(reply)
 
+def restrict(bot, update):
+    if "group" not in update.message.chat.type:
+        update.message.reply_text("That only works in groups.")
+        return
+    ident = str(update.message.chat.id)
+    if not ident in chatlogs:
+        chat = update.message.chat
+        title = get_chatname(chat)
+        chatlog = Chatlog(chat.id, chat.type, title)
+        chatlogs[chatlog.id] = chatlog
+    else:
+        chatlog = chatlogs[ident]
+    if chatlog.is_restricted():
+        user = update.message.chat.get_member(update.message.from_user.id)
+        if not user_is_admin(user):
+            reply = "You do not have permissions to do that."
+            update.message.reply_text(reply)
+            return
+    chatlogs[ident].toggle_restrict()
+    reply = (chatlogs[ident].is_restricted() and "I will only let admins " or "I will let everyone ") + "configure me now."
+    update.message.reply_text(reply)
+
 def stop(bot, update):
     global ADMIN_ID
     chatlog = chatlogs[update.message.chat.id]
@@ -261,6 +310,7 @@ def stop(bot, update):
     print("I got blocked by user " + chatlog.id)
 
 def main():
+    global ADMIN_ID
     parser = argparse.ArgumentParser(description='A Telegram markov bot.')
     parser.add_argument('token', metavar='TOKEN', help='The Bot Token to work with the Telegram Bot API')
     parser.add_argument('admin_id', metavar='ADMIN_ID', type=int, help='The ID of the Telegram user that manages this bot')
@@ -290,6 +340,7 @@ def main():
     dp.add_handler(CommandHandler("stop", stop))
     dp.add_handler(CommandHandler("speak", speak))
     dp.add_handler(CommandHandler("answer", set_answer_freq))
+    dp.add_handler(CommandHandler("restrict", restrict))
 
     # on noncommand i.e message - echo the message on Telegram
     # dp.add_handler(MessageHandler(Filters.text, echo))