Feature: Sqlite Database Checking and Addition

main.py

@@ -5,6 +5,7 @@
 '''
 from imaplib import IMAP4_SSL
 from mail_extract import get_plugin_name
+import util_db
 import email
 import queue
 import threading
@@ -25,10 +26,10 @@ class myMailThread(threading.Thread):
         self.login()
 
     def run(self):
-        print("Starting " + self.name)
+        # print("Starting " + self.name)
         process_queue(self.name, self.mail, self.q)
         self.logout()
-        print("Exiting " + self.name)
+        # print("Exiting " + self.name)
 
     def login(self):
         self.mail.login(self.user, self.password)
@@ -77,7 +78,9 @@ def process_queue(threadName, mail, q):
         if not plugin:
             print("%s - Failed - %s " % (threadName, subject))
             continue
-        print("%s - Processed - %s" % (threadName, plugin))
+        conn = util_db.db_sqlite()
+        conn.add_plugin(plugin, "plugin")
+        # print("%s - Processed - %s" % (threadName, plugin))
 
 
 if __name__ == '__main__':
@@ -96,18 +99,18 @@ if __name__ == '__main__':
     password = "password"
     totalVulns = {}
 
+    # Build our singular query thread and update our workQueue
+    thr = myMailThread(threadID, nameList, workQueue, user, password)
+    thr.get_messages()
+    threadID += 1
+
     # Start all our processing threads
     for x in range(0, 3):
         tName = "{}{}".format(parserName, x)
         thread = myMailThread(threadID, tName, workQueue, user, password)
         thread.start()
         threads.append(thread)
-        threadID += 1
-
-    # Build our singular query thread and update our workQueue
-    thr = myMailThread(threadID, nameList, workQueue, user, password)
-    thr.get_messages()
-    threadID += 1
+        threadID += 2
 
     # Wait until the queue is empty to quit
     while not workQueue.empty():

util_db.py

@@ -14,36 +14,44 @@ class db_sqlite():
         try:
             self.conn.execute("create table plugin(name, version)")
         except Exception:
-            print("Table: plugin already exists")
+            pass
+            # print("Table: plugin already exists")
         try:
             self.conn.execute("create table monitor(name, version)")
         except Exception:
-            print("Table: watch_list already exists")
-        print("Database schema for " + db_file + " at v0.1")
+            pass
+            # print("Table: watch_list already exists")
+        # print("Database schema for " + db_file + " at v0.1")
 
     def add_plugin(self, plugin, table):
         try:
             name = plugin["plugin"]
             version = plugin["version"]
-            if not self.check_plugin(name, version, table):
+            if self.check_plugin(name, version, table):
                 temp_plugin = (name, version)
                 self.conn.execute('INSERT INTO {} VALUES (?,?)'.format(table), temp_plugin)
                 self.conn.commit()
-                print("Plugin: " + name + ", " + version + "Inserted in " + table)
+                print("Inserted: Plugin: " + name + ", " + version + " in table: " + table)
             else:
-                print("Plugin: " + name + ", " + version + "Already in " + table)
+                print("Already Exists: " + name + ", " + version + " in table: " + table)
         except Exception:
             raise NameError
 
-        print(name, version)
-
     def check_plugin(self, plugin_name, plugin_version, table):
-        plugin = (plugin_name, plugin_version)
-        data = self.conn.execute('SELECT * FROM {} WHERE name=? AND version=?'.format(table), plugin)
-        if data:
-            return (True, "plugin exists")
-        else:
-            return (False, "not in database")
+        """
+        Checks for plugin name and version existance
+        returns: True if it does not exist yet
+        """
+        temp_plugin = (plugin_name, plugin_version)
+        c = self.conn.cursor()
+        c.execute('SELECT * FROM {} WHERE name=? AND version=?'.format(table), temp_plugin)
+        data = c.fetchone()
+        # Do the biggest hack on attempting to view a row of something that should not exist
+        try:
+            for row in data:
+                return False
+        except Exception:
+            return True
 
     def cleanup(self):
         self.conn.commit()