Python: Magic -> Special and reaarange classes

Author: yoff

python/ql/src/experimental/dataflow/internal/DataFlowPrivate.qll

@@ -1,6 +1,6 @@
 private import python
 private import DataFlowPublic
-import semmle.python.Magic
+import semmle.python.SpecialMethods
 
 //--------
 // Data flow graph
@@ -160,7 +160,7 @@ class DataFlowClassValue extends DataFlowCallable, TClassValue {
 
 newtype TDataFlowCall =
   TCallNode(CallNode call) or
-  TMagicCall(MagicMethod::Actual magic)
+  TSpecialCall(SpecialMethodCallNode special)
 
 abstract class DataFlowCall extends TDataFlowCall {
   /** Gets a textual representation of this element. */
@@ -179,7 +179,6 @@ abstract class DataFlowCall extends TDataFlowCall {
   abstract DataFlowCallable getEnclosingCallable();
 }
 
-
 /** Represents a call to a callable. */
 class CallNodeCall extends DataFlowCall, TCallNode {
   CallNode call;
@@ -192,9 +191,7 @@ class CallNodeCall extends DataFlowCall, TCallNode {
 
   override string toString() { result = call.toString() }
 
-  override ControlFlowNode getArg(int n) {
-    result = call.getArg(n)
-  }
+  override ControlFlowNode getArg(int n) { result = call.getArg(n) }
 
   override ControlFlowNode getNode() { result = call }
 
@@ -203,22 +200,24 @@ class CallNodeCall extends DataFlowCall, TCallNode {
   override DataFlowCallable getEnclosingCallable() { result.getScope() = call.getNode().getScope() }
 }
 
-class MagicCall extends DataFlowCall, TMagicCall {
-  MagicMethod::Actual magic;
+class SpecialCall extends DataFlowCall, TSpecialCall {
+  SpecialMethodCallNode special;
 
-  MagicCall() { this = TMagicCall(magic) }
+  SpecialCall() { this = TSpecialCall(special) }
 
-  override string toString() { result = magic.toString() }
+  override string toString() { result = special.toString() }
 
-  override ControlFlowNode getArg(int n) {
-    result = magic.(MagicMethod::Potential).getArg(n)
-  }
+  override ControlFlowNode getArg(int n) { result = special.(SpecialMethod::Potential).getArg(n) }
 
-  override ControlFlowNode getNode() { result = magic }
+  override ControlFlowNode getNode() { result = special }
 
-  override DataFlowCallable getCallable() { result = TCallableValue(magic.getResolvedMagicMethod()) }
+  override DataFlowCallable getCallable() {
+    result = TCallableValue(special.getResolvedSpecialMethod())
+  }
 
-  override DataFlowCallable getEnclosingCallable() { result.getScope() = magic.getNode().getScope() }
+  override DataFlowCallable getEnclosingCallable() {
+    result.getScope() = special.getNode().getScope()
+  }
 }
 
 /** A data flow node that represents a call argument. */

python/ql/src/semmle/python/Magic.qll

@@ -0,0 +1,117 @@
+/**
+ * Provides support for special methods.
+ * This is done in two steps:
+ *   - A subset of `ControlFlowNode`s are labelled as potentially corresponding to
+ *     a special method call (by being an instance of `SpecialMethod::Potential`).
+ *   - A subset of the potential special method calls are labelled as being actual
+ *     special method calls (`SpecialMethodCallNode`) if the appropriate method is defined.
+ * Extend `SpecialMethod::Potential` to capture more cases.
+ */
+
+import python
+
+
+/** A control flow node which might correpsond to a special method call. */
+class PotentialSpecialMethodCallNode extends ControlFlowNode {
+  PotentialSpecialMethodCallNode() { this instanceof SpecialMethod::Potential}
+}
+
+/**
+ * Machinery for detecting special method calls.
+ * Extend `SpecialMethod::Potential` to capture more cases.
+ */
+module SpecialMethod {
+  /** A control flow node which might correpsond to a special method call. */
+  abstract class Potential extends ControlFlowNode {
+    /** Gets the name of the method that would be called. */
+    abstract string getSpecialMethodName();
+
+    /** Gets the controlflow node that would be passed as the specified argument. */
+    abstract ControlFlowNode getArg(int n);
+
+    /**
+     * Gets the control flow node corresponding to the instance
+     * that would define the special method.
+     */
+    ControlFlowNode getSelf() { result = this.getArg(0) }
+  }
+
+  /** A binary expression node that might correspond to a special method call. */
+  class SpecialBinOp extends Potential, BinaryExprNode {
+    Operator operator;
+
+    SpecialBinOp() { this.getOp() = operator }
+
+    override string getSpecialMethodName() { result = operator.getSpecialMethodName() }
+
+    override ControlFlowNode getArg(int n) {
+      n = 0 and result = this.getLeft()
+      or
+      n = 1 and result = this.getRight()
+    }
+  }
+
+  /** A subscript expression node that might correspond to a special method call. */
+  abstract class SpecialSubscript extends Potential, SubscriptNode {
+    override ControlFlowNode getArg(int n) {
+      n = 0 and result = this.getObject()
+      or
+      n = 1 and result = this.getIndex()
+    }
+  }
+
+  /** A subscript expression node that might correspond to a call to __getitem__. */
+  class SpecialGetItem extends SpecialSubscript {
+    SpecialGetItem() { this.isLoad() }
+
+    override string getSpecialMethodName() { result = "__getitem__" }
+  }
+
+  /** A subscript expression node that might correspond to a call to __setitem__. */
+  class SpecialSetItem extends SpecialSubscript {
+    SpecialSetItem() { this.isStore() }
+
+    override string getSpecialMethodName() { result = "__setitem__" }
+
+    override ControlFlowNode getArg(int n) {
+      n = 0 and result = this.getObject()
+      or
+      n = 1 and result = this.getIndex()
+      or
+      n = 2 and result = this.getValueNode()
+    }
+
+    private ControlFlowNode getValueNode() {
+      exists(AssignStmt a |
+        a.getATarget() = this.getNode() and
+        result.getNode() = a.getValue()
+      )
+      or
+      exists(AugAssign a |
+        a.getTarget() = this.getNode() and
+        result.getNode() = a.getValue()
+      )
+    }
+  }
+
+  /** A subscript expression node that might correspond to a call to __delitem__. */
+  class SpecialDelItem extends SpecialSubscript {
+    SpecialDelItem() { this.isDelete() }
+
+    override string getSpecialMethodName() { result = "__delitem__" }
+  }
+}
+
+class SpecialMethodCallNode extends PotentialSpecialMethodCallNode {
+  Value resolvedSpecialMethod;
+
+  SpecialMethodCallNode() {
+    exists(SpecialMethod::Potential pot |
+      this.(SpecialMethod::Potential) = pot and
+      pot.getSelf().pointsTo().getClass().lookup(pot.getSpecialMethodName()) = resolvedSpecialMethod
+    )
+  }
+
+  /** The method that is called. */
+  Value getResolvedSpecialMethod() { result = resolvedSpecialMethod }
+}