Python: Add API graph support for subclasses

tausbn · tausbn · commit ef600575cace · 2021-02-05T16:52:58.000+01:00
diff --git a/python/ql/src/semmle/python/ApiGraphs.qll b/python/ql/src/semmle/python/ApiGraphs.qll
@@ -92,6 +92,11 @@ module API {
      */
     Node getReturn() { result = getASuccessor(Label::return()) }
 
+    /**
+     * Gets a node representing a subclass of the class represented by this node.
+     */
+    Node getASubclass() { result = getASuccessor(Label::subclass()) }
+
     /**
      * Gets a string representation of the lexicographically least among all shortest access paths
      * from the root to this node.
@@ -358,6 +363,12 @@ module API {
         // Calling a node that is a use of `base`
         lbl = Label::return() and
         ref = pred.getACall()
+        or
+        // Subclassing a node
+        lbl = Label::subclass() and
+        exists(DataFlow::Node superclass | pred.flowsTo(superclass) |
+          ref.asExpr().(ClassExpr).getABase() = superclass.asExpr()
+        )
       )
     }
 
@@ -468,4 +479,6 @@ private module Label {
 
   /** Gets the `return` edge label. */
   string return() { result = "getReturn()" }
+
+  string subclass() { result = "getASubclass()" }
 }
diff --git a/python/ql/test/experimental/dataflow/ApiGraphs/test.py b/python/ql/test/experimental/dataflow/ApiGraphs/test.py
@@ -79,3 +79,20 @@ def f():
 from unknown import * #$ use=moduleImport("unknown")
 
 hello() #$ MISSING: use=moduleImport("unknown").getMember("hello").getReturn()
+
+
+# Subclasses
+
+from flask.views import View #$ use=moduleImport("flask").getMember("views").getMember("View")
+
+class MyView(View): #$ use=moduleImport("flask").getMember("views").getMember("View").getASubclass()
+    pass
+
+instance = MyView() #$ use=moduleImport("flask").getMember("views").getMember("View").getASubclass().getReturn()
+
+def internal():
+    from pflask.views import View #$ use=moduleImport("pflask").getMember("views").getMember("View")
+    class IntMyView(View): #$ use=moduleImport("pflask").getMember("views").getMember("View").getASubclass()
+        pass
+
+    int_instance = IntMyView() #$ use=moduleImport("pflask").getMember("views").getMember("View").getASubclass().getReturn()
