Python: Model invoke.context.Context

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

@@ -34,7 +34,7 @@ private module Invoke {
    * WARNING: Only holds for a few predefined attributes.
    */
   private DataFlow::Node invoke_attr(DataFlow::TypeTracker t, string attr_name) {
-    attr_name in ["run", "sudo"] and
+    attr_name in ["run", "sudo", "context", "Context"] and
     (
       t.start() and
       result = DataFlow::importMember("invoke", attr_name)
@@ -69,15 +69,73 @@ private module Invoke {
   }
 
   /** Provides models for the `invoke` module. */
-  module invoke { }
+  module invoke {
+    /** Gets a reference to the `invoke.context` module. */
+    DataFlow::Node context() { result = invoke_attr("context") }
+
+    /** Provides models for the `invoke.context` module */
+    module context {
+      /** Provides models for the `invoke.context.Context` class */
+      module Context {
+        /** Gets a reference to the `invoke.context.Context` class. */
+        private DataFlow::Node class_(DataFlow::TypeTracker t) {
+          t.start() and
+          result = DataFlow::importMember("invoke.context", "Context")
+          or
+          t.startInAttr("Context") and
+          result = invoke::context()
+          or
+          t.start() and
+          result = invoke_attr("Context")
+          or
+          exists(DataFlow::TypeTracker t2 | result = class_(t2).track(t2, t))
+        }
+
+        /** Gets a reference to the `invoke.context.Context` class. */
+        DataFlow::Node class_() { result = class_(DataFlow::TypeTracker::end()) }
+
+        /** Gets a reference to an instance of `invoke.context.Context`. */
+        private DataFlow::Node instance(DataFlow::TypeTracker t) {
+          t.start() and
+          result.asCfgNode().(CallNode).getFunction() =
+            invoke::context::Context::class_().asCfgNode()
+          or
+          exists(DataFlow::TypeTracker t2 | result = instance(t2).track(t2, t))
+        }
+
+        /** Gets a reference to an instance of `invoke.context.Context`. */
+        DataFlow::Node instance() { result = instance(DataFlow::TypeTracker::end()) }
+
+        /** Gets a reference to the `run` or `sudo` methods on a `invoke.context.Context` instance. */
+        private DataFlow::Node instanceRunMethods(DataFlow::TypeTracker t) {
+          t.startInAttr(["run", "sudo"]) and
+          result = invoke::context::Context::instance()
+          or
+          exists(DataFlow::TypeTracker t2 | result = instanceRunMethods(t2).track(t2, t))
+        }
+
+        /** Gets a reference to the `run` or `sudo` methods on a `invoke.context.Context` instance. */
+        DataFlow::Node instanceRunMethods() {
+          result = instanceRunMethods(DataFlow::TypeTracker::end())
+        }
+      }
+    }
+  }
 
   /**
-   * A call to either of the `invoke.run` or `invoke.sudo` functions
-   * See http://docs.pyinvoke.org/en/stable/api/__init__.html
+   * A call to either
+   * - `invoke.run` or `invoke.sudo` functions (http://docs.pyinvoke.org/en/stable/api/__init__.html)
+   * - `run` or `sudo` methods on a `invoke.context.Context` instance (http://docs.pyinvoke.org/en/stable/api/context.html#invoke.context.Context.run)
    */
   private class InvokeRunCommandCall extends SystemCommandExecution::Range {
     InvokeRunCommandCall() {
-      this.asCfgNode().(CallNode).getFunction() = invoke_attr(["run", "sudo"]).asCfgNode()
+      exists(DataFlow::Node callFunction |
+        this.asCfgNode().(CallNode).getFunction() = callFunction.asCfgNode()
+      |
+        callFunction = invoke_attr(["run", "sudo"])
+        or
+        callFunction = invoke::context::Context::instanceRunMethods()
+      )
     }
 
     override DataFlow::Node getCommand() {