Running queries

aes.py

@@ -0,0 +1,29 @@
+import base64
+import hashlib
+from Crypto import Random
+from Crypto.Cipher import AES
+
+class AESCipher(object):
+
+    def __init__(self, key): 
+        self.bs = 32
+        self.key = hashlib.sha256(key.encode()).digest()
+
+    def encrypt(self, raw):
+        raw = self._pad(raw)
+        iv = Random.new().read(AES.block_size)
+        cipher = AES.new(self.key, AES.MODE_CBC, iv)
+        return base64.b64encode(iv + cipher.encrypt(raw))
+
+    def decrypt(self, enc):
+        enc = base64.b64decode(enc)
+        iv = enc[:AES.block_size]
+        cipher = AES.new(self.key, AES.MODE_CBC, iv)
+        return self._unpad(cipher.decrypt(enc[AES.block_size:])).decode('utf-8')
+
+    def _pad(self, s):
+        return s + (self.bs - len(s) % self.bs) * chr(self.bs - len(s) % self.bs)
+
+    @staticmethod
+    def _unpad(s):
+        return s[:-ord(s[len(s)-1:])]

app.py

@@ -1,7 +1,8 @@
-from flask import Flask, render_template, request, redirect, url_for, flash, jsonify
+from flask import Flask, render_template, request, redirect, url_for, flash
 
 import db
 import config
+from aes import AESCipher
 
 app = Flask(__name__)
 app.secret_key = config.flask_secret_key
@@ -16,12 +17,15 @@ def index():
                 queries = db.get_queries()
         else:
             queries = db.get_queries()
+
+        databases = db.get_databases()
 
-        return render_template("index.html", queries=queries, issearchword=searchword)
+        return render_template("index.html", databases=databases, queries=queries, issearchword=searchword)
 
 @app.route("/queries.json/")
 def query_list():
     queries = db.get_queries()
+    databases = db.get_databases()
     qlist = [
                 {
                     "id":str(query["_id"]), 
@@ -35,13 +39,13 @@ def query_list():
 @app.route("/query/", methods=["POST"])
 def query_add():
     try:
-        title = request.form["title"].strip()
-        sql = request.form["sql"].strip()
-        tags = request.form["tags"].strip()
-        desc = request.form["desc"].strip()
-        who = request.form["who"].strip()
+        title = request.form.get("title").strip()
+        sql = request.form.get("sql").strip()
+        tags = request.form.get("tags").strip()
+        desc = request.form.get("desc").strip()
+        who = request.form.get("who").strip()
 
-        if len(title) == 0 or len(sql) == 0:
+        if not title or not sql:
             flash("Title and Sql are required fields", "error")
             return redirect(url_for("index"))
 
@@ -56,8 +60,9 @@ def query_add():
 
 @app.route("/query/<id>/", methods=["GET", "DELETE"])
 def query_view(id):
+    databases = db.get_databases()
     query = db.get_query_details(id)
-    return render_template("view_query.html", query=query)
+    return render_template("view_query.html", databases=databases, query=query)
 
 @app.route("/query/<id>/json/", methods=["GET"])
 def query_json_view(id):
@@ -108,5 +113,139 @@ def query_delete(id):
         flash("Delete Successful", "success")
         return redirect(url_for("index"))
 
+@app.route("/query/<id>/run/<database_id>", methods=["GET", "POST"])
+def query_run(id, database_id):
+    """ Run a query against a specific database instance """
+    if config.enable_run_query:
+        database = db.get_database_details(database_id)
+        query = db.get_query_details(id)
+        query_results = None
+        query_results_cols = []
+        error = None
+
+        # try and import the DB engine
+        try: 
+            dbapi2 = __import__(config.database_engine)
+        except ImportError as e:
+            app.logger.error("Fatal error. Could not import DB engine.", exc_info=True)
+            flash("Fatal error. Contact Administrator", "error")
+            return redirect(url_for("index"))
+
+        # try and make the connection and run the query
+        try:
+            if database.get("password"):
+                crypt = AESCipher(config.flask_secret_key)
+                password = crypt.decrypt(database.get("password"))
+            else:
+                password = None
+
+            connect = dbapi2.connect(database=database.get("name"), 
+                host=database.get("hostname"), 
+                port=database.get("port"), 
+                user=database.get("user"), 
+                password=password)
+
+            curse = connect.cursor()
+            curse.execute(query["sql"])
+            query_results = curse.fetchall()
+
+            # Assemble column names so the table makes sense
+            for col in curse.description:
+                query_results_cols.append(col.name)
+
+        except dbapi2.ProgrammingError, e:
+            # TODO: Exceptions don't seem to be standard in DB-API2, 
+            # so this will likely have to be checked against other
+            # engines.  The following works with psycopg2.
+            if hasattr(e, "pgerror"):
+                error = e.pgerror
+            else:
+                error = "There was an error with your query."
+        except dbapi2.Error as e:
+            if hasattr(e, "pgerror"):
+                error = e.pgerror or e.message
+                app.logger.error(error)
+            else:
+                error = e.msg
+                app.logger.error(error)
+
+    else: 
+        database = None
+        query = None
+        query_results = None
+        query_results_cols = None
+        error = None
+    return render_template("run_query.html", run_enabled=config.enable_run_query, query=query, database=database, query_results=query_results, query_results_cols=query_results_cols, error=error)
+
+@app.route("/database/", methods=["POST"])
+def database_add():
+    try:
+        name = request.form.get("name").strip()
+        hostname = request.form.get("hostname").strip()
+        port = request.form.get("port", "").strip()
+        user = request.form.get("user", "").strip()
+        password = request.form.get("password", "").strip()
+        desc = request.form.get("desc", "").strip()
+
+        if not name or not hostname:
+            flash("Name and Host/File Name are required fields", "error")
+            return redirect(url_for("index"))
+
+        db.insert_database(name, hostname, port, user, password, desc)
+        flash("Database Added!", "success")
+        return redirect(url_for("index"))
+
+    except Exception as e:
+        app.logger.error("Fatal error", exc_info=True)
+        flash("Fatal error. Contact Administrator", "error")
+        return redirect(url_for("index"))
+
+@app.route("/database/<id>/edit/", methods=["GET", "POST"])
+def database_edit(id):
+    if request.method == "GET":
+        database = db.get_database_details(id)
+        if database['password']:
+            database['password'] = 'placeholder'
+        return render_template("edit_database.html", database=database)
+    elif request.method == "POST":
+        try:
+            id = request.form["id"].strip()
+            name = request.form.get("name").strip()
+            hostname = request.form.get("hostname").strip()
+            port = request.form.get("port", "").strip()
+            user = request.form.get("user", "").strip()
+            password = request.form.get("password", "").strip()
+            if password == 'placeholder':
+                password = None
+            desc = request.form.get("desc", "").strip()
+
+            if not name or not hostname:
+                flash("Name and Host/File Name are required fields", "error")
+                return redirect(url_for("database_edit", id=id))
+
+            db.update_database(id, name, hostname, port, user, password, desc)
+            flash("Database Modified!", "success")
+            return redirect(url_for("database_view", id=id))
+
+        except Exception as e:
+            app.logger.error("Fatal error", exc_info=True)
+            flash("Fatal error. Contact Administrator", "error")
+            return redirect(url_for("index"))
+
+@app.route("/database/<id>/", methods=["GET", "DELETE"])
+def database_view(id):
+    database = db.get_database_details(id)
+    return render_template("view_database.html", database=database)
+
+@app.route("/database/<id>/delete/", methods=["GET", "POST"])
+def database_delete(id):
+    if request.method == "GET":
+        database = db.get_database_details(id)
+        return render_template("delete_database.html", database=database)
+    elif request.method == "POST":
+        db.delete_database(id)
+        flash("Delete Successful", "success")
+        return redirect(url_for("index"))
+
 if __name__ == "__main__":
     app.run(debug=config.flask_debug, port=config.flask_port, host=config.flask_bind_address)

config.py.sample

@@ -2,12 +2,21 @@
 flask_debug=True
 flask_bind_address="0.0.0.0"
 flask_port=5000
+# For security purposes, this field should be sufficiently randomized
 flask_secret_key=""
 
 # Mongo properties
 mongo_hostname="localhost"
 mongo_port=27017
 mongo_db="company_queries"
 mongo_collection="queries"
+mongo_collection_database="databases"
 mongo_username=""
 mongo_password=""
+
+# Whether or not running queries is allowed
+enable_run_query = False
+
+# The database engine should be DB-API2 compatible module that can be
+# imported.  More info: https://wiki.python.org/moin/DbApiFaq
+database_engine = 'psycopg2'

db.py

@@ -2,6 +2,7 @@
 from bson.objectid import ObjectId
 
 import config
+from aes import AESCipher
 
 client = m.MongoClient(config.mongo_hostname, config.mongo_port)
 db = client[config.mongo_db]
@@ -11,6 +12,7 @@
     db.authenticate(config.mongo_username, config.mongo_password)
 
 queries = db[config.mongo_collection]
+databases = db[config.mongo_collection_database]
 
 def get_tags_list(tags):
     return [ tag.strip() for tag in tags.strip().split(",") \
@@ -61,3 +63,50 @@ def get_queries():
 
 def get_query_details(id):
     return queries.find_one(ObjectId(id))
+
+def insert_database(name, hostname, port=None, user=None, password=None, desc=None):
+    if password:
+        crypt = AESCipher(config.flask_secret_key)
+        encrypted_password = crypt.encrypt(password)
+    else:
+        encrypted_password = None
+
+    databases.insert({
+            "name": name,
+            "hostname": hostname,
+            "port": port,
+            "user": user, 
+            "password": encrypted_password,
+            "desc": desc
+        })
+
+def update_database(id, name, hostname, port=None, user=None, password=None, desc=None):
+
+    db_set = {
+        "name": name,
+        "hostname": hostname,
+        "port": port,
+        "user": user, 
+        "desc": desc
+    }
+
+    if password:
+        crypt = AESCipher(config.flask_secret_key)
+        encrypted_password = crypt.encrypt(password)
+        db_set["password"] = encrypted_password
+    else:
+        encrypted_password = None
+
+    databases.update({
+        "_id": ObjectId(id)
+        },
+        { "$set" : db_set}, upsert=False)
+
+def delete_database(id):
+    databases.remove({"_id": ObjectId(id)})
+
+def get_databases():
+    return list(databases.find())
+
+def get_database_details(id):
+    return databases.find_one(ObjectId(id))

requirements.txt

@@ -6,3 +6,4 @@ itsdangerous==0.24
 nose==1.3.6
 pymongo==2.8
 wsgiref==0.1.2
+pycrypto>=2.6.1

templates/delete_database.html

@@ -0,0 +1,50 @@
+{% extends "base.html" %}
+
+{% block pagescripts %}
+
+  <link rel="stylesheet" href="/static/highlightjs/styles/railscasts.css">
+  <script src="/static/highlightjs/highlight.pack.js"></script>
+  <script>hljs.initHighlightingOnLoad();</script>
+
+  <style>
+    pre, code {
+      background: #232323;
+      border:none;
+      font-family: monospace;
+    }
+  </style>
+
+{% endblock %}
+
+{% block container %}
+
+  <div class="row">
+    <div class="col-md-8 col-md-offset-2">
+  <div class="panel panel-default">
+    <div class="panel-heading">
+      <h3 class="panel-title">Are you sure to delete this Database?</h3>
+    </div>
+    <div class="panel-body">
+      <form method="POST" action="{{ url_for("database_delete", id=database._id) }}">
+
+        <div class="row">
+          <div class="col-md-12">
+            <h1>{{ database.name }}</h1>
+          </div>
+        </div>
+
+        <hr/>
+
+        <div class="row">
+          <div class="col-md-1">
+            <button class="btn btn-danger" type="submit">Yes, Delete it!</button>
+          </div>
+        </div>
+      </form>
+    </div>
+  </div>
+    </div>
+  </div>
+
+
+{% endblock %}