Python: Make builtin compile function additional taint step

Author: RasmusWL

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

@@ -347,7 +347,7 @@ private module Stdlib {
    * WARNING: Only holds for a few predefined attributes.
    */
   private DataFlow::Node builtins_attr(DataFlow::TypeTracker t, string attr_name) {
-    attr_name in ["exec", "eval"] and
+    attr_name in ["exec", "eval", "compile"] and
     (
       t.start() and
       result = DataFlow::importMember(["builtins", "__builtin__"], attr_name)
@@ -419,6 +419,21 @@ private module Stdlib {
 
     override DataFlow::Node getCode() { result.asCfgNode() = call.getArg(0) }
   }
+
+  /** An additional taint step for calls to the builtin function `compile` */
+  private class BuiltinsCompileCallAdditionalTaintStep extends TaintTracking::AdditionalTaintStep {
+    override predicate step(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
+      exists(CallNode call |
+        nodeTo.asCfgNode() = call and
+        call.getFunction() = builtins_attr("compile").asCfgNode() and
+        (
+          call.getArg(0) = nodeFrom.asCfgNode()
+          or
+          call.getArgByName("source") = nodeFrom.asCfgNode()
+        )
+      )
+    }
+  }
 }
 
 /**

python/ql/test/experimental/library-tests/frameworks/stdlib/TestTaint.expected

@@ -1,4 +1,4 @@
 | CodeExecution.py:35 | ok   | test_additional_taint | src |
-| CodeExecution.py:36 | fail | test_additional_taint | cmd1 |
-| CodeExecution.py:37 | fail | test_additional_taint | cmd2 |
-| CodeExecution.py:38 | fail | test_additional_taint | cmd3 |
+| CodeExecution.py:36 | ok   | test_additional_taint | cmd1 |
+| CodeExecution.py:37 | ok   | test_additional_taint | cmd2 |
+| CodeExecution.py:38 | ok   | test_additional_taint | cmd3 |

python/ql/test/experimental/query-tests/Security-new-dataflow/CWE-094/CodeInjection.expected

@@ -1,10 +1,13 @@
 edges
 | code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | code_injection.py:7:10:7:13 | ControlFlowNode for code |
 | code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | code_injection.py:8:10:8:13 | ControlFlowNode for code |
+| code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | code_injection.py:10:10:10:12 | ControlFlowNode for cmd |
 nodes
 | code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | semmle.label | ControlFlowNode for Attribute |
 | code_injection.py:7:10:7:13 | ControlFlowNode for code | semmle.label | ControlFlowNode for code |
 | code_injection.py:8:10:8:13 | ControlFlowNode for code | semmle.label | ControlFlowNode for code |
+| code_injection.py:10:10:10:12 | ControlFlowNode for cmd | semmle.label | ControlFlowNode for cmd |
 #select
 | code_injection.py:7:10:7:13 | ControlFlowNode for code | code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | code_injection.py:7:10:7:13 | ControlFlowNode for code | $@ flows to here and is interpreted as code. | code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | A user-provided value |
 | code_injection.py:8:10:8:13 | ControlFlowNode for code | code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | code_injection.py:8:10:8:13 | ControlFlowNode for code | $@ flows to here and is interpreted as code. | code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | A user-provided value |
+| code_injection.py:10:10:10:12 | ControlFlowNode for cmd | code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | code_injection.py:10:10:10:12 | ControlFlowNode for cmd | $@ flows to here and is interpreted as code. | code_injection.py:6:12:6:23 | ControlFlowNode for Attribute | A user-provided value |