PyMongo and Mongoengine sinks

python/ql/src/experimental/semmle/python/frameworks/Stdlib.qll

@@ -11,9 +11,60 @@ private import experimental.semmle.python.Concepts
 private import semmle.python.ApiGraphs
 
 /**
- * PyMongoQuery
- * MongoEngineQuery
+ * PyMongoCall
+ * MongoEngineCall
  * Custom escapes
  */
-private module NoSQL { }
-// introduce more json libs like python\ql\src\semmle\python\frameworks\Stdlib.qll:941
+private module NoSQL {
+  // doesn't work currently
+  private class PyMongoCall extends DataFlow::Node, NoSQLQuery::Range {
+    DataFlow::Node queryNode;
+
+    PyMongoCall() {
+      exists(SsaVariable clientVar, CallNode findCall |
+        (
+          clientVar.getDefinition().getImmediateDominator() =
+            Value::named("pymongo.MongoClient").getACall() or
+          clientVar.getDefinition().getImmediateDominator() =
+            Value::named("flask_pymongo.PyMongo").getACall()
+        ) and
+        clientVar.getAUse().getNode() = findCall.getNode().getFunc().(Attribute).getObject() and
+        findCall.getNode().getFunc().(Attribute).getName().matches("%find%") and
+        this.asCfgNode() = findCall and
+        queryNode.asExpr() = findCall.getArg(0).getNode()
+      )
+    }
+
+    override DataFlow::Node getQueryNode() { result = queryNode }
+  }
+
+  // `API::moduleImport("mongoengine").getMember("Document").getASubclass*().getACall()` doesn't point
+  // to our sinks
+  private class MongoEngineCall extends DataFlow::CallCfgNode, NoSQLQuery::Range {
+    DataFlow::Node queryNode;
+
+    MongoEngineCall() {
+      exists(DataFlow::AttrRead objectsMethod |
+        this.getFunction() = objectsMethod and
+        API::moduleImport("mongoengine").getMember("Document").getASubclass*().getACall() =
+          objectsMethod.getObject().getALocalSource() and
+        queryNode = this.getArg(0)
+      )
+    }
+
+    override DataFlow::Node getQueryNode() { result = queryNode }
+  }
+
+  // pending: look for more Sanitizer libs
+  private class MongoSanitizer extends DataFlow::CallCfgNode, NoSQLSanitizer::Range {
+    DataFlow::Node escapeNode;
+
+    MongoSanitizer() {
+      this =
+        API::moduleImport("mongosanitizer").getMember("sanitizer").getMember("sanitize").getACall() and
+      escapeNode = this.getArg(0)
+    }
+
+    override DataFlow::Node getSanitizerNode() { result = escapeNode }
+  }
+}