Merge pull request #1472 from markshannon/python-taint-through-iterators

Python: Track taint through iteration and iterators including generators.

python/ql/src/semmle/python/security/TaintTracking.qll

@@ -159,6 +159,12 @@ abstract class TaintKind extends string {
 
     string repr() { result = this }
 
+    /** Gets the taint resulting from iterating over this kind of taint.
+     * For example iterating over a text file produces lines. So iterating
+     * over a tainted file would result in tainted strings
+     */
+    TaintKind getTaintForIteration() { none() }
+
 }
 
 /** Taint kinds representing collections of other taint kind.
@@ -212,6 +218,10 @@ class SequenceKind extends CollectionKind {
         result = "sequence of " + itemKind
     }
 
+    override TaintKind getTaintForIteration() {
+        result = itemKind
+    }
+
 }
 
 
@@ -863,6 +873,10 @@ library module TaintFlowImplementation {
         or
         call_taint_step(fromnode, totaint, tocontext, tonode)
         or
+        iteration_step(fromnode, totaint, tocontext, tonode)
+        or
+        yield_step(fromnode, totaint, tocontext, tonode)
+        or
         exists(DataFlowNode fromnodenode |
             fromnodenode = fromnode.getNode() and
             (
@@ -1053,6 +1067,26 @@ library module TaintFlowImplementation {
         )
     }
 
+    predicate yield_step(TaintedNode fromnode, TrackedValue totaint, CallContext tocontext, CallNode call) {
+        exists(PyFunctionObject func |
+            func.getFunction().isGenerator() and
+            func.getACall() = call and
+            (
+                fromnode.getContext() = tocontext.getCallee(call)
+                or
+                fromnode.getContext() = tocontext and tocontext = TTop()
+            ) and
+            exists(Yield yield |
+                yield.getScope() = func.getFunction() and
+                yield.getValue() = fromnode.getNode().getNode()
+            ) and
+            exists(SequenceKind seq |
+                seq.getItem() = fromnode.getTaintKind() and
+                totaint = fromnode.getTrackedValue().toKind(seq)
+            )
+        )
+    }
+
     predicate call_taint_step(TaintedNode fromnode, TrackedValue totaint, CallContext tocontext, CallNode call) {
         exists(string name |
             call.getFunction().(AttrNode).getObject(name) = fromnode.getNode() and
@@ -1067,6 +1101,13 @@ library module TaintFlowImplementation {
         )
     }
 
+    /** Holds if `v` is defined by a `for` statement, the definition being `defn` */
+    cached predicate iteration_step(TaintedNode fromnode, TrackedValue totaint, CallContext tocontext, ControlFlowNode iter) {
+        exists(ForNode for | for.iterates(iter, fromnode.getNode())) and
+        totaint = TTrackedTaint(fromnode.getTaintKind().getTaintForIteration()) and
+        tocontext = fromnode.getContext()
+    }
+
     predicate self_init_end_transfer(EssaVariable self, CallContext callee, CallNode call, CallContext caller) {
         exists(ClassValue cls, Function init |
             call.getFunction().pointsTo(cls) and
@@ -1161,6 +1202,9 @@ library module TaintFlowImplementation {
         tainted_with(def, context, origin)
         or
         tainted_exception_capture(def, context, origin)
+        or
+        tainted_iteration(def, context, origin)
+
     }
 
     predicate tainted_scope_entry(ScopeEntryDefinition def, CallContext context, TaintedNode origin) {
@@ -1363,6 +1407,12 @@ library module TaintFlowImplementation {
         context = fromnode.getContext()
     }
 
+    pragma [noinline]
+    private predicate tainted_iteration(IterationDefinition def, CallContext context, TaintedNode fromnode) {
+        def.getDefiningNode() = fromnode.getNode() and
+        context = fromnode.getContext()
+    }
+
     /* A call that returns a copy (or similar) of the argument */
     predicate copyCall(ControlFlowNode fromnode, CallNode tonode) {
         tonode.getFunction().(AttrNode).getObject("copy") = fromnode