Merge pull request github#6064 from tausbn/python-add-get-method-call

Python: Add `getAMethodCall` to `LocalSourceNode`

Author: RasmusWL

python/change-notes/2021-06-15-add-method-call-conveniences.md

@@ -0,0 +1,5 @@
+lgtm,codescanning
+* A new class `DataFlow::MethodCallNode` extends `DataFlow::CallCfgNode` with convenient methods for
+  accessing the receiver and method name of a method call.
+* The `LocalSourceNode` class now has a `getAMethodCall` method, with which one can easily access
+  method calls with the given node as a receiver.

python/ql/src/Security/CWE-327/Ssl.qll

@@ -39,12 +39,10 @@ API::Node sslContextInstance() {
   result = API::moduleImport("ssl").getMember(["SSLContext", "create_default_context"]).getReturn()
 }
 
-class WrapSocketCall extends ConnectionCreation, DataFlow::CallCfgNode {
+class WrapSocketCall extends ConnectionCreation, DataFlow::MethodCallNode {
   WrapSocketCall() { this = sslContextInstance().getMember("wrap_socket").getACall() }
 
-  override DataFlow::Node getContext() {
-    result = this.getFunction().(DataFlow::AttrRead).getObject()
-  }
+  override DataFlow::Node getContext() { result = this.getObject() }
 }
 
 class OptionsAugOr extends ProtocolRestriction, DataFlow::CfgNode {

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

@@ -64,15 +64,14 @@ private module Re {
    *
    * See https://docs.python.org/3/library/re.html#regular-expression-objects
    */
-  private class CompiledRegex extends DataFlow::CallCfgNode, RegexExecution::Range {
+  private class CompiledRegex extends DataFlow::MethodCallNode, RegexExecution::Range {
     DataFlow::Node regexNode;
 
     CompiledRegex() {
-      exists(DataFlow::CallCfgNode patternCall, DataFlow::AttrRead reMethod |
-        this.getFunction() = reMethod and
+      exists(DataFlow::MethodCallNode patternCall |
         patternCall = API::moduleImport("re").getMember("compile").getACall() and
-        patternCall.flowsTo(reMethod.getObject()) and
-        reMethod.getAttributeName() instanceof RegexExecutionMethods and
+        patternCall.flowsTo(this.getObject()) and
+        this.getMethodName() instanceof RegexExecutionMethods and
         regexNode = patternCall.getArg(0)
       )
     }

python/ql/src/semmle/python/dataflow/new/internal/Attributes.qll

@@ -191,7 +191,7 @@ private class ClassDefinitionAsAttrWrite extends AttrWrite, CfgNode {
  * - Dynamic attribute reads using `getattr`: `getattr(object, attr)`
  * - Qualified imports: `from module import attr as name`
  */
-abstract class AttrRead extends AttrRef, Node { }
+abstract class AttrRead extends AttrRef, Node, LocalSourceNode { }
 
 /** A simple attribute read, e.g. `object.attr` */
 private class AttributeReadAsAttrRead extends AttrRead, CfgNode {

python/ql/src/semmle/python/dataflow/new/internal/DataFlowPublic.qll

@@ -180,6 +180,45 @@ 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(...)`. Note that this is only tracked through local scope.
+ */
+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.
+   *
+   * Note that this method may have multiple results if a single call node represents calls to
+   * multiple different objects and methods. If you want to link up objects and method names
+   * accurately, use the `calls` method instead.
+   */
+  string getMethodName() { result = method_lookup.getAttributeName() }
+
+  /**
+   * Gets the data-flow node corresponding to the object receiving this call. That is, the `foo` in
+   * `foo.bar(...)`.
+   *
+   * Note that this method may have multiple results if a single call node represents calls to
+   * multiple different objects and methods. If you want to link up objects and method names
+   * accurately, use the `calls` method instead.
+   */
+  Node getObject() { result = method_lookup.getObject() }
+
+  /** Holds if this data-flow node calls method `methodName` on the object node `object`. */
+  predicate calls(Node object, string methodName) {
+    // As `getObject` and `getMethodName` may both have multiple results, we must look up the object
+    // and method name directly on `method_lookup`.
+    object = method_lookup.getObject() and
+    methodName = method_lookup.getAttributeName()
+  }
+}
+
 /**
  * An expression, viewed as a node in a data flow graph.
  *

python/ql/src/semmle/python/dataflow/new/internal/LocalSources.qll

@@ -78,6 +78,16 @@ class LocalSourceNode extends Node {
    */
   CallCfgNode getACall() { Cached::call(this, result) }
 
+  /**
+   * Gets a call to the method `methodName` on this node.
+   *
+   * Includes both calls that have the syntactic shape of a method call (as in `obj.m(...)`), and
+   * calls where the callee undergoes some additional local data flow (as in `tmp = obj.m; m(...)`).
+   */
+  MethodCallNode getAMethodCall(string methodName) {
+    result = this.getAnAttributeRead(methodName).getACall()
+  }
+
   /**
    * Gets a node that this node may flow to using one heap and/or interprocedural step.
    *

python/ql/src/semmle/python/frameworks/Cryptography.qll

@@ -228,11 +228,7 @@ private module CryptographyModel {
     /** Gets a reference to the encryptor of a Cipher instance using algorithm with `algorithmName`. */
     DataFlow::LocalSourceNode cipherEncryptor(DataFlow::TypeTracker t, string algorithmName) {
       t.start() and
-      exists(DataFlow::AttrRead attr |
-        result.(DataFlow::CallCfgNode).getFunction() = attr and
-        attr.getAttributeName() = "encryptor" and
-        attr.getObject() = cipherInstance(algorithmName)
-      )
+      result.(DataFlow::MethodCallNode).calls(cipherInstance(algorithmName), "encryptor")
       or
       exists(DataFlow::TypeTracker t2 | result = cipherEncryptor(t2, algorithmName).track(t2, t))
     }
@@ -249,11 +245,7 @@ private module CryptographyModel {
     /** Gets a reference to the dncryptor of a Cipher instance using algorithm with `algorithmName`. */
     DataFlow::LocalSourceNode cipherDecryptor(DataFlow::TypeTracker t, string algorithmName) {
       t.start() and
-      exists(DataFlow::AttrRead attr |
-        result.(DataFlow::CallCfgNode).getFunction() = attr and
-        attr.getAttributeName() = "decryptor" and
-        attr.getObject() = cipherInstance(algorithmName)
-      )
+      result.(DataFlow::MethodCallNode).calls(cipherInstance(algorithmName), "decryptor")
       or
       exists(DataFlow::TypeTracker t2 | result = cipherDecryptor(t2, algorithmName).track(t2, t))
     }
@@ -271,18 +263,14 @@ private module CryptographyModel {
      * An encrypt or decrypt operation from `cryptography.hazmat.primitives.ciphers`.
      */
     class CryptographyGenericCipherOperation extends Cryptography::CryptographicOperation::Range,
-      DataFlow::CallCfgNode {
+      DataFlow::MethodCallNode {
       string algorithmName;
 
       CryptographyGenericCipherOperation() {
-        exists(DataFlow::AttrRead attr |
-          this.getFunction() = attr and
-          attr.getAttributeName() = ["update", "update_into"] and
-          (
-            attr.getObject() = cipherEncryptor(algorithmName)
-            or
-            attr.getObject() = cipherDecryptor(algorithmName)
-          )
+        exists(DataFlow::Node object, string method |
+          object in [cipherEncryptor(algorithmName), cipherDecryptor(algorithmName)] and
+          method in ["update", "update_into"] and
+          this.calls(object, method)
         )
       }
 
@@ -337,16 +325,10 @@ private module CryptographyModel {
      * An hashing operation from `cryptography.hazmat.primitives.hashes`.
      */
     class CryptographyGenericHashOperation extends Cryptography::CryptographicOperation::Range,
-      DataFlow::CallCfgNode {
+      DataFlow::MethodCallNode {
       string algorithmName;
 
-      CryptographyGenericHashOperation() {
-        exists(DataFlow::AttrRead attr |
-          this.getFunction() = attr and
-          attr.getAttributeName() = "update" and
-          attr.getObject() = hashInstance(algorithmName)
-        )
-      }
+      CryptographyGenericHashOperation() { this.calls(hashInstance(algorithmName), "update") }
 
       override Cryptography::CryptographicAlgorithm getAlgorithm() {
         result.matchesName(algorithmName)

python/ql/test/experimental/dataflow/method-calls/test.expected

@@ -0,0 +1,8 @@
+conjunctive_lookup
+| test.py:6:1:6:6 | ControlFlowNode for meth() | meth() | obj1 | bar |
+| test.py:6:1:6:6 | ControlFlowNode for meth() | meth() | obj1 | foo |
+| test.py:6:1:6:6 | ControlFlowNode for meth() | meth() | obj2 | bar |
+| test.py:6:1:6:6 | ControlFlowNode for meth() | meth() | obj2 | foo |
+calls_lookup
+| test.py:6:1:6:6 | ControlFlowNode for meth() | meth() | obj1 | foo |
+| test.py:6:1:6:6 | ControlFlowNode for meth() | meth() | obj2 | bar |

python/ql/test/experimental/dataflow/method-calls/test.py

@@ -0,0 +1,6 @@
+if cond:
+    meth = obj1.foo
+else:
+    meth = obj2.bar
+
+meth()

python/ql/test/experimental/dataflow/method-calls/test.ql

@@ -0,0 +1,18 @@
+import python
+import semmle.python.dataflow.new.DataFlow
+import experimental.dataflow.TestUtil.PrintNode
+
+query predicate conjunctive_lookup(
+  DataFlow::MethodCallNode methCall, string call, string object, string methodName
+) {
+  call = prettyNode(methCall) and
+  object = prettyNode(methCall.getObject()) and
+  methodName = methCall.getMethodName()
+}
+
+query predicate calls_lookup(
+  DataFlow::MethodCallNode methCall, string call, string object, string methodName
+) {
+  call = prettyNode(methCall) and
+  exists(DataFlow::Node o | methCall.calls(o, methodName) and object = prettyNode(o))
+}