Python: Add MethodCallNode class

tausbn · web-flow · commit d4b05547ba3d · 2021-06-15T15:04:29.000Z
Roughly patterned after the JS equivalent.
diff --git a/python/ql/src/semmle/python/dataflow/new/internal/DataFlowPublic.qll b/python/ql/src/semmle/python/dataflow/new/internal/DataFlowPublic.qll
@@ -180,6 +180,30 @@ class CallCfgNode extends CfgNode, LocalSourceNode {
   Node getArgByName(string name) { result.asCfgNode() = node.getArgByName(name) }
 }
 
+/**
+ * A data-flow node corresponding to a method call, that is `foo.bar(...)`.
+ *
+ * Also covers the case where the method lookup is done separately from the call itself, as in
+ * `temp = foo.bar; temp(...)`.
+ */
+class MethodCallNode extends CallCfgNode {
+  AttrRead method_lookup;
+
+  MethodCallNode() { method_lookup = this.getFunction().getALocalSource() }
+
+  /** Gets the name of the method being invoked (the `bar` in `foo.bar(...)`, if it can be determined. */
+  string getMethodName() { result = method_lookup.getAttributeName() }
+
+  /** Gets the data-flow node corresponding to the receiver of this call. That is, the `foo` in `foo.bar(...)`. */
+  Node getReceiver() { result = method_lookup.getObject() }
+
+  /** Holds if this data-flow node calls method `methodName` on receiver node `receiver`. */
+  predicate calls(Node receiver, string methodName) {
+    receiver = this.getReceiver() and
+    methodName = this.getMethodName()
+  }
+}
+
 /**
  * An expression, viewed as a node in a data flow graph.
  *
