plugin-alert/util_db.py

#!/usr/bin/env python

import sqlite3
db_file = "plugin-vuln.db"


class db_sqlite():
    def __init__(self):
        self.db_file = db_file
        self.conn = sqlite3.connect(db_file)
        self.create_db()

    def create_db(self):
        try:
            self.conn.execute("create table plugin(name, version)")
        except Exception:
            pass
            # print("Table: plugin already exists")
        try:
            self.conn.execute("create table monitor(name, version)")
        except Exception:
            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 self.check_plugin(name, version, table):
                temp_plugin = (name, version)
                self.conn.execute('INSERT INTO {} VALUES (?,?)'.format(table), temp_plugin)
                self.conn.commit()
                print("Inserted: Plugin: " + name + ", " + version + " in table: " + table)
            else:
                print("Already Exists: " + name + ", " + version + " in table: " + table)
        except Exception:
            raise NameError

    def check_plugin(self, plugin_name, plugin_version, table):
        """
        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()
        self.conn.close()