Add additional query tests

To ensure that this query works against numerous usages of libraries such as PyMongo, Flask PyMongo, Mongoengine, and Flask Mongoengine, I've added a variety of query tests to test against. These tests deal with scenarious such as:

- Subscript expressions
- Mongoengine instances and Document subclasses
- Mongoengine connection usage
- And more...

python/ql/test/experimental/query-tests/Security/CWE-943/flask_mongoengine_get_db_bad.py

@@ -0,0 +1,31 @@
+from flask import Flask, request
+from flask_mongoengine import MongoEngine
+import json
+
+app = Flask(__name__)
+db = MongoEngine(app)
+db.init_app(app)
+
+
+class Movie(db.Document):
+    title = db.StringField(required=True)
+    year = db.IntField()
+    rated = db.StringField()
+    director = db.StringField()
+    actors = db.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsanitized_search = request.args['search']
+    json_search = json.loads(unsanitized_search)
+
+    retrieved_db = db.get_db()
+    result = retrieved_db["Movie"].find_one_or_404({'name': json_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/flask_mongoengine_get_db_good.py

@@ -0,0 +1,33 @@
+from flask import Flask, request
+from flask_mongoengine import MongoEngine
+from mongosanitizer.sanitizer import sanitize
+import json
+
+app = Flask(__name__)
+db = MongoEngine(app)
+db.init_app(app)
+
+
+class Movie(db.Document):
+    title = db.StringField(required=True)
+    year = db.IntField()
+    rated = db.StringField()
+    director = db.StringField()
+    actors = db.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsanitized_search = request.args['search']
+    json_search = json.loads(unsanitized_search)
+    safe_search = sanitize(json_search)
+
+    retrieved_db = db.get_db()
+    result = retrieved_db["Movie"].find_one_or_404({'name': safe_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoclient_subscript_bad.py

@@ -0,0 +1,29 @@
+from flask import Flask, request
+import mongoengine as me
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+
+    db = me.connect('mydb')  
+    data = db['mydb']['movie'].find({'name': json_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoclient_subscript_good.py

@@ -0,0 +1,31 @@
+from flask import Flask, request
+import mongoengine as me
+from mongosanitizer.sanitizer import sanitize
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+    safe_search = sanitize(json_search)
+
+    db = me.connect('mydb')
+    data = db['mydb']['movie'].find({'name': safe_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine-flask-db-document-subclass-bad.py

@@ -0,0 +1,29 @@
+from flask import Flask, request
+from flask_mongoengine import MongoEngine
+import json
+
+app = Flask(__name__)
+app.config.from_pyfile('the-config.cfg')
+db = MongoEngine(app)
+
+class Movie(db.Document):
+    title = db.StringField(required=True)
+    year = db.IntField()
+    rated = db.StringField()
+    director = db.StringField()
+    actors = db.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+
+    data = Movie.objects(__raw__=json_search)
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine-flask-db-document-subclass-good.py

@@ -0,0 +1,31 @@
+from flask import Flask, request
+from flask_mongoengine import MongoEngine
+from mongosanitizer.sanitizer import sanitize
+import json
+
+app = Flask(__name__)
+app.config.from_pyfile('the-config.cfg')
+db = MongoEngine(app)
+
+class Movie(db.Document):
+    title = db.StringField(required=True)
+    year = db.IntField()
+    rated = db.StringField()
+    director = db.StringField()
+    actors = db.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+    safe_search = sanitize(json_search)
+
+    data = Movie.objects(__raw__=safe_search)
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine_connect_bad.py

@@ -0,0 +1,29 @@
+from flask import Flask, request
+import mongoengine as me
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+
+    db = me.connect('mydb')  
+    data = db.mydb.movie.find({'name': json_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine_connect_good.py

@@ -0,0 +1,31 @@
+from flask import Flask, request
+import mongoengine as me
+from mongosanitizer.sanitizer import sanitize
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+    safe_search = sanitize(json_search)
+
+    db = me.connect('mydb')  
+    data = db.mydb.movie.find({'name': safe_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine_connect_via_connection_bad.py

@@ -0,0 +1,30 @@
+from flask import Flask, request
+import mongoengine as me
+from mongoengine.connection import get_db, connect
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+
+    db = connect('mydb')
+    data = db.mydb.movie.find({'name': json_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine_connect_via_connection_good.py

@@ -0,0 +1,32 @@
+from flask import Flask, request
+import mongoengine as me
+from mongoengine.connection import get_db, connect
+from mongosanitizer.sanitizer import sanitize
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+    safe_search = sanitize(json_search)
+
+    db = connect('mydb')
+    data = db.mydb.movie.find({'name': json_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine_get_db_bad.py

@@ -0,0 +1,29 @@
+from flask import Flask, request
+import mongoengine as me
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+
+    db = me.get_db()
+    data = db.movie.find({'name': json_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine_get_db_good.py

@@ -0,0 +1,31 @@
+from flask import Flask, request
+import mongoengine as me
+from mongosanitizer.sanitizer import sanitize
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+    safe_search = sanitize(json_search)
+
+    db = me.get_db()
+    data = db.movie.find({'name': safe_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine_get_db_via_connection_bad.py

@@ -0,0 +1,30 @@
+from flask import Flask, request
+import mongoengine as me
+from mongoengine.connection import get_db, connect
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+
+    db = get_db()
+    data = db.movie.find({'name': json_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)

python/ql/test/experimental/query-tests/Security/CWE-943/mongoengine_get_db_via_connection_good.py

@@ -0,0 +1,32 @@
+from flask import Flask, request
+import mongoengine as me
+from mongoengine.connection import get_db, connect
+from mongosanitizer.sanitizer import sanitize
+import json
+
+app = Flask(__name__)
+
+
+class Movie(me.Document):
+    title = me.StringField(required=True)
+    year = me.IntField()
+    rated = me.StringField()
+    director = me.StringField()
+    actors = me.ListField()
+
+
+Movie(title='aa').save()
+Movie(title='bb').save()
+
+
+@app.route("/")
+def home_page():
+    unsafe_search = request.args['search']
+    json_search = json.loads(unsafe_search)
+    safe_search = sanitize(json_search)
+
+    db = get_db()
+    data = db.movie.find({'name': safe_search})
+
+# if __name__ == "__main__":
+#     app.run(debug=True)