Refactor, remove queue usage, and add debug mode

Author: KevinMGranger

rabbit/main.py

@@ -1,53 +1,27 @@
 import sys
-if sys.version_info < (3,0,0):
-    print("Please run me in Python 3.")
-    sys.exit(0)
-
-from .stackoverflowchatsession import StackOverflowChatSession
-import config
-
-import asyncio
+import os
 import json
 import html
+import logging
 import random
-from queue import Queue
+import threading
+import re
 from .dbmodel import User, get_session
+from .sochat import StackOverflowChatSession, EventType
+
+import config
 
 PYTHON_ROOM_ID = 6
 PERSONAL_SANDBOX_ROOM_ID = 118024
 ROTATING_KNIVES_ROOM_ID = 71097
+AUTHORIZED_USERS = {
+            953482, #Kevin
+            6621329 #Terry
+        }
 
-#room that all of the bot's actions will occur in.
-#todo: make it possible for the bot to operate simultaneously in multiple rooms. The chat API supports this natively, I just didn't account for it in my design.
 PRIMARY_ROOM_ID = PYTHON_ROOM_ID
 
-event_type_names = [
-    "placeholder because ids are 1-indexed",
-    "message posted",
-    "message edited",
-    "user entered",
-    "user left",
-    "room name changed",
-    "message starred",
-    "UNKNOWN",
-    "user mentioned",
-    "message flagged",
-    "message deleted",
-    "file added",
-    "moderator flag",
-    "user settings chagned",
-    "global notification",
-    "account level changed",
-    "user notification",
-    "invitation",
-    "message reply",
-    "message moved out",
-    "message moved in",
-    "time break",
-    "feed ticker",
-    "user suspended",
-    "user merged",
-]
+logger = logging.getLogger('rabbit')
 
 def abbreviate(msg, maxlen=25):
     if len(msg) < maxlen: return msg
@@ -67,13 +41,11 @@ class Rabbit(StackOverflowChatSession):
             - "kick [user id]" - kicks the user, if bot has RO rights
             - "move [message id,message id,message id]" - moves one or more messages to the Rotating Knives room, if bot has RO rights
     """
-    def __init__(self, email, password, admin_message_queue): 
+    def __init__(self, email, password, room, trash_room, authorized_users):
         StackOverflowChatSession.__init__(self, email, password)
-        self.admin_message_queue = admin_message_queue
-        self.authorized_users = {
-            953482, #Kevin
-            6621329 #Terry
-        }
+        self.room = room
+        self.trash_room = trash_room
+        self.authorized_users = authorized_users
 
     def onConnect(self, response):
         print('Connected:', response.peer)
@@ -83,20 +55,18 @@ def onOpen(self):
 
     def onMessage(self, payload):
         d = json.loads(payload.decode("utf-8"))
-        for ROOM_ID, data in d.items():
+        logger.debug("Payload: {}".format(d))
+        for room_id, data in d.items():
             if "e" not in data: #some kind of keepalive message that we don't care about
                 continue
             for event in data["e"]:
-                event_type = event["event_type"]
-                if event_type >= len(event_type_names):
-                    raise Exception("Unrecognized event type: {} \nWith event data:".format(event_type, event))
+                try:
+                    event_type = event["event_type"]
+                    event_type = EventType(event_type)
+                except ValueError:
+                    raise Exception("Unrecognized event type: {} \nWith event data: {}".format(event_type, event))
                 if event_type == 1: #ordinary user message
-                    content = html.unescape(event["content"])
-                    print(abbreviate("{}: {}".format(event["user_name"], content), 119))
-                    if event["user_id"] in self.authorized_users: #possible administrator command
-                        if content == "!ping":
-                            print("Detected a command. Replying...")
-                            self.send_message(PRIMARY_ROOM_ID, "pong")
+                    self._on_regular_message(event)
                 elif event_type in (3,4): #user entered/left
                     action = {3:"entered", 4:"left"}[event_type]
                     print("user {} {} room {}".format(repr(event["user_name"]), action, repr(event["room_name"])))
@@ -109,62 +79,57 @@ def onMessage(self, payload):
 
                         #now post a picture of a bunny.
                         bunny_url = random.choice(config.kick_reply_images)
-                        self.send_message(PRIMARY_ROOM_ID, bunny_url)
+                        self.send_message(self.room, bunny_url)
                     else:
                         print("Info: Unknown event content {} in account level changed event.".format(repr(event["content"])))
                         print(event)
                 else:
-                    print(event_type_names[event_type])
+                    logger.info("Event: {}".format(event_type))
 
-    def onClose(self, was_clean, code, reason):
-          print('Closed:', reason)
-          import sys; sys.exit(0)
+    def _on_regular_message(self, event):
+        # TODO: never ever react to self messages
+        content = html.unescape(event["content"])
+        print(abbreviate("{}: {}".format(event["user_name"], content), 119))
+        if event["user_id"] in self.authorized_users: #possible administrator command
+            if content == "!ping":
+                print("Detected a command. Replying...")
+                self.send_message(self.room, "pong")
 
-    def onIdle(self):
-        while not self.admin_message_queue.empty():
-            msg = self.admin_message_queue.get()
-            self.onAdminMessage(msg)
+    def onClose(self, was_clean, code, reason):
+        print('Closed:', reason)
+        sys.exit(0)
 
     def onAdminMessage(self, msg):
         print("Got admin message: {}".format(msg))
         if msg == "shutdown":
             print("Shutting down...")
-            import sys; sys.exit(0)
+            sys.exit(0)
         elif msg.startswith("say"):
-            self.send_message(PRIMARY_ROOM_ID, msg.partition(" ")[2])
+            self.send_message(self.room, msg.partition(" ")[2])
         elif msg.startswith("cancel"):
             messageId = msg.partition(" ")[2]
             self.cancel_stars(messageId)
         elif msg.startswith("kick"):
             userId = msg.partition(" ")[2]
-            self.kick(PRIMARY_ROOM_ID, userId)
+            self.kick(self.room, userId)
         elif msg.startswith("move"):
             messageIds = msg.partition(" ")[2].split()
-            self.move_messages(PRIMARY_ROOM_ID, messageIds, ROTATING_KNIVES_ROOM_ID)
+            self.move_messages(self.room, messageIds, self.trash_room)
         else:
             print("Sorry, didn't understand that command.")
 
 
-#connect to the SO chat server. This function never returns.
-def create_and_run_chat_session(admin_message_queue = None):
-    if admin_message_queue is None:
-        admin_message_queue = Queue()
-
-    session = Rabbit(config.email, config.password, admin_message_queue)
-    session.join_and_run_forever(PRIMARY_ROOM_ID)
-
-import threading
 
 #create a GUI the user can use to send admin commands. This function never returns.
 #(hint: use threads if you want to run both this and `create_and_run_chat_session`)
-def create_admin_window(message_queue):
+def create_admin_window(bot):
     from tkinter import Tk, Entry, Button
 
     def clicked():
-        message_queue.put(box.get())
+        bot.loop.call_soon_threadsafe(bot.onAdminMessage, box.get())
 
     def on_closing():
-        message_queue.put("shutdown")
+        bot.loop.call_soon_threadsafe(bot.onAdminMessage, "shutdown")
         root.destroy()
 
     root = Tk()
@@ -179,8 +144,23 @@ def on_closing():
 
 
 def main():
-    message_queue = Queue()
-    t = threading.Thread(target=create_admin_window, args=(message_queue,))
+    session = Rabbit(config.email, config.password, PRIMARY_ROOM_ID, ROTATING_KNIVES_ROOM_ID, AUTHORIZED_USERS)
+    t = threading.Thread(target=create_admin_window, args=(session,))
     t.start()
 
-    create_and_run_chat_session(message_queue)
+    session.join_and_run_forever(PRIMARY_ROOM_ID)
+
+
+def debug():
+    handler = logging.StreamHandler()
+    handler.setLevel(logging.DEBUG)
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)
+
+    session = Rabbit(config.email, config.password,
+        os.environ['room'], os.environ['trash'], set(int(x) for x in os.environ['users'].split(':')))
+
+    if 'tk' in os.environ:
+        t = threading.Thread(target=create_admin_window, args=(session,))
+        t.start()
+    session.join_and_run_forever(os.environ['room'])

rabbit/stackoverflowchatsession.py renamed to rabbit/sochat.py

@@ -1,19 +1,52 @@
 import requests
 import json
 import asyncio
+import logging
+import enum
 from urllib.parse import quote_plus
 from bs4 import BeautifulSoup as BS
 from autobahn.asyncio.websocket import WebSocketClientProtocol, WebSocketClientFactory
 
 
+logger = logging.getLogger('rabbit')
+
+EventType = enum.IntEnum('EventType', [
+    "message_posted",
+    "message_edited",
+    "user_entered",
+    "user_left",
+    "room_name_changed",
+    "message_starred",
+    "UNKNOWN",
+    "user_mentioned",
+    "message_flagged",
+    "message_deleted",
+    "file_added",
+    "moderator_flag",
+    "user_settings_chagned",
+    "global_notification",
+    "account_level_changed",
+    "user_notification",
+    "invitation",
+    "message_reply",
+    "message_moved_out",
+    "message_moved_in",
+    "time_break",
+    "feed_ticker",
+    "user_suspended",
+    "user_merged",
+])
+
 class StackOverflowChatSession:
     def __init__(self, email, password):
         url = "https://stackoverflow.com/users/login"
         login_data = {"email": email, "password": password}
         session = requests.Session()
+        logger.debug("Logging in")
         session.post(url,login_data)
         #TODO: perform some cursory checking to confirm that logging in actually worked
 
+        logger.debug("Getting cookie")
         x = session.get("http://chat.stackoverflow.com")
 
         soup = BS(x.content, "html.parser")
@@ -71,8 +104,6 @@ def onMessage(self, payload):
         pass
     def onClose(self, was_clean, code, reason):
         pass
-    def onIdle(self):
-        pass
 
     def join_and_run_forever(self, roomid):
         session = self
@@ -88,16 +119,11 @@ def onClose(self, was_clean, code, reason): session.onClose(was_clean, code, rea
         factory = WebSocketClientFactory(url, headers={"Origin":"http://chat.stackoverflow.com"})
         factory.protocol = SoClient
         self.loop = asyncio.get_event_loop()
-        self.loop.call_later(1, self._onIdle)
         coro = self.loop.create_connection(factory, host, 80)
         self.loop.run_until_complete(coro)
         self.loop.run_forever()
         self.loop.close()
 
-    def _onIdle(self):
-        self.onIdle()
-        self.loop.call_later(1, self._onIdle)
-
     def _get_webservice_url(self, roomid):
         x = self._post("http://chat.stackoverflow.com/ws-auth", {"roomid":roomid})
 
@@ -111,6 +137,7 @@ def _post(self, url, params=None):
         Send a POST message using some predetermined headers and params that are always necessary when talking to SO.
         using this instead of requests.post will save you the effort of adding the fkey, cookie, etc yourself every time.
         """
+        logger.debug("Posting to {}".format(url))
         if params is None:
             params = {}
         params["fkey"] = self.fkey

setup.py

@@ -19,6 +19,7 @@
         console_scripts=[
             'rabbit_userscript=rabbit.userscript_server:main',
             'rabbit=rabbit.main:main',
+            'bugs=rabbit.main:debug'
         ]
     )
 )