Adds user whitelisting support

db_logger.py

@@ -6,11 +6,13 @@
 class DatabaseLogger:
     conn = None
     cursor = None
+    senders_whitelist = None
 
-    def __init__(self, host, name, user, password):
+    def __init__(self, host, name, user, password, senders_whitelist=None):
         self.conn = psycopg2.connect(host=host, dbname=name, user=user, password=password)
         self.conn.autocommit = True
         self.cursor = self.conn.cursor()
+        self.senders_whitelist = senders_whitelist
 
     def close(self):
         self.cursor.close()
@@ -24,7 +26,11 @@ def log_chat(self, sender, message, channel):
             return
 
         try:
-            self.cursor.execute("INSERT INTO chat_log (sender, message, channel, date) VALUES (%s, %s, %s, %s)",
+            if self.senders_whitelist is not None and sender in self.senders_whitelist:
+                self.cursor.execute("INSERT INTO chat_log (sender, message, channel, date) VALUES (%s, %s, %s, %s)",
+                                (sender, message, channel, current_time_in_milli()))
+            elif self.senders_whitelist is None:
+                self.cursor.execute("INSERT INTO chat_log (sender, message, channel, date) VALUES (%s, %s, %s, %s)",
                                 (sender, message, channel, current_time_in_milli()))
         except psycopg2.DataError as e:
             print e

manager.py

@@ -17,12 +17,19 @@ class TwitchManager:
     SECONDS_BETWEEN_UPDATE_STREAMS = 60
     SECONDS_BETWEEN_CREATE_BOTS = 15
 
-    def __init__(self, channels_amount, channels, log_filename=None):
+    def __init__(self, channels_amount, channels, log_filename=None, users_whitelist=None):
         self.bots = []
         self.channels_amount = channels_amount
         self.log_filename = log_filename
         self.channels = channels
-        self.db_logger = DatabaseLogger(settings.DATABASE['HOST'],
+        if settings.USERS is not None and len(settings.USERS) is not 0:
+            self.db_logger = DatabaseLogger(settings.DATABASE['HOST'],
+                                        settings.DATABASE['NAME'],
+                                        settings.DATABASE['USER'],
+                                        settings.DATABASE['PASSWORD'],
+                                        settings.USERS)
+        else:
+            self.db_logger = DatabaseLogger(settings.DATABASE['HOST'],
                                         settings.DATABASE['NAME'],
                                         settings.DATABASE['USER'],
                                         settings.DATABASE['PASSWORD'])
@@ -33,10 +40,18 @@ def _create_bot(self, name, channels):
                              settings.IRC['NICK'],
                              settings.IRC['PASSWORD'],
                              self.log_filename)
-        bot_db_logger = DatabaseLogger(settings.DATABASE['HOST'],
+        if settings.USERS is not None and len(settings.USERS) is not 0:
+            bot_db_logger = DatabaseLogger(settings.DATABASE['HOST'],
                                        settings.DATABASE['NAME'],
                                        settings.DATABASE['USER'],
-                                       settings.DATABASE['PASSWORD'])
+                                       settings.DATABASE['PASSWORD'],
+                                       settings.USERS)
+        else:
+            bot_db_logger = DatabaseLogger(settings.DATABASE['HOST'],
+	                               settings.DATABASE['NAME'],
+	                               settings.DATABASE['USER'],
+				       settings.DATABASE['PASSWORD'])
+
         bot = TwitchBot(name, conn, bot_db_logger, Queue.Queue(), self.log_filename)
         bot.daemon = True
         bot.connect_and_join_channels(channels)

settings.py.example

@@ -16,4 +16,11 @@ DATABASE = {
 
 API = {
 	'CLIENTID': 'YOUR_CLIENT_ID'
-}
+}
+
+# Uncomment to only log certain users.
+#USERS = [
+#    'yourusername',
+#    'friendsname'
+#]
+