Python: Summaries for container store step

python/ql/lib/semmle/python/dataflow/new/FlowSummary.qll

@@ -28,6 +28,27 @@ module SummaryComponent {
   /** Gets a summary component that represents a list element. */
   SummaryComponent listElement() { result = content(any(ListElementContent c)) }
 
+  /** Gets a summary component that represents a set element. */
+  SummaryComponent setElement() { result = content(any(SetElementContent c)) }
+
+  /** Gets a summary component that represents a tuple element. */
+  SummaryComponent tupleElement(int index) {
+    exists(TupleElementContent c | c.getIndex() = index and result = content(c))
+  }
+
+  /** Gets a summary component that represents a dictionary element. */
+  SummaryComponent dictionaryElement(string key) {
+    exists(DictionaryElementContent c | c.getKey() = key and result = content(c))
+  }
+
+  /** Gets a summary component that represents a dictionary element at any key. */
+  SummaryComponent dictionaryElementAny() { result = content(any(DictionaryElementAnyContent c)) }
+
+  /** Gets a summary component that represents an attribute element. */
+  SummaryComponent attribute(string attr) {
+    exists(AttributeContent c | c.getAttribute() = attr and result = content(c))
+  }
+
   /** Gets a summary component that represents the return value of a call. */
   SummaryComponent return() { result = SC::return(any(ReturnKind rk)) }
 }

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll

@@ -765,8 +765,8 @@ predicate readStep(Node nodeFrom, Content c, Node nodeTo) {
   or
   matchReadStep(nodeFrom, c, nodeTo)
   or
-  popReadStep(nodeFrom, c, nodeTo)
-  or
+  // popReadStep(nodeFrom, c, nodeTo)
+  // or
   forReadStep(nodeFrom, c, nodeTo)
   or
   attributeReadStep(nodeFrom, c, nodeTo)

python/ql/lib/semmle/python/dataflow/new/internal/FlowSummaryImplSpecific.qll

@@ -105,6 +105,27 @@ predicate neutralElement(FlowSummary::SummarizedCallable c, string provenance) {
 SummaryComponent interpretComponentSpecific(AccessPathToken c) {
   c = "ListElement" and
   result = FlowSummary::SummaryComponent::listElement()
+  or
+  c = "SetElement" and
+  result = FlowSummary::SummaryComponent::setElement()
+  or
+  exists(int index |
+    c.getAnArgument("TupleElement") = index.toString() and
+    result = FlowSummary::SummaryComponent::tupleElement(index)
+  )
+  or
+  exists(string key |
+    c.getAnArgument("DictionaryElement") = key and
+    result = FlowSummary::SummaryComponent::dictionaryElement(key)
+  )
+  or
+  c = "DictionaryElementAny" and
+  result = FlowSummary::SummaryComponent::dictionaryElementAny()
+  or
+  exists(string attr |
+    c.getAnArgument("Attribute") = attr and
+    result = FlowSummary::SummaryComponent::attribute(attr)
+  )
 }
 
 /** Gets the textual representation of a summary component in the format used for flow summaries. */

python/ql/lib/semmle/python/dataflow/new/internal/TaintTrackingPrivate.qll

@@ -1,6 +1,7 @@
 private import python
 private import semmle.python.dataflow.new.DataFlow
 private import semmle.python.dataflow.new.internal.DataFlowPrivate as DataFlowPrivate
+private import FlowSummaryImpl as FlowSummaryImpl
 private import semmle.python.dataflow.new.internal.TaintTrackingPublic
 private import semmle.python.ApiGraphs
 
@@ -55,6 +56,8 @@ private module Cached {
     awaitStep(nodeFrom, nodeTo)
     or
     asyncWithStep(nodeFrom, nodeTo)
+    or
+    FlowSummaryImpl::Private::Steps::summaryLocalStep(nodeFrom, nodeTo, false)
   }
 }
 
@@ -159,7 +162,7 @@ predicate stringManipulation(DataFlow::CfgNode nodeFrom, DataFlow::CfgNode nodeT
  * is currently very imprecise, as an example, since we model `dict.get`, we treat any
  * `<tainted object>.get(<arg>)` will be tainted, whether it's true or not.
  */
-predicate containerStep(DataFlow::CfgNode nodeFrom, DataFlow::Node nodeTo) {
+predicate containerStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
   // construction by literal
   //
   // TODO: once we have proper flow-summary modeling, we might not need this step any
@@ -181,36 +184,36 @@ predicate containerStep(DataFlow::CfgNode nodeFrom, DataFlow::Node nodeTo) {
   // don't provide that right now.
   DataFlowPrivate::comprehensionStoreStep(nodeFrom, _, nodeTo)
   or
-  // constructor call
-  exists(DataFlow::CallCfgNode call | call = nodeTo |
-    call = API::builtin(["list", "set", "frozenset", "dict", "tuple"]).getACall() and
-    call.getArg(0) = nodeFrom
-    // TODO: Properly handle defaultdict/namedtuple
-  )
-  or
-  // functions operating on collections
-  exists(DataFlow::CallCfgNode call | call = nodeTo |
-    call = API::builtin(["sorted", "reversed", "iter", "next"]).getACall() and
-    call.getArg(0) = nodeFrom
-  )
-  or
-  // methods
-  exists(DataFlow::MethodCallNode call, string methodName | call = nodeTo |
-    methodName in [
-        // general
-        "copy", "pop",
-        // dict
-        "values", "items", "get", "popitem"
-      ] and
-    call.calls(nodeFrom, methodName)
-  )
-  or
+  // // methods
+  // exists(DataFlow::MethodCallNode call, string methodName | call = nodeTo |
+  //   methodName in [
+  //       // general
+  //       "copy", "pop",
+  //       // dict
+  //       "values", "items", "get", "popitem"
+  //     ] and
+  //   call.calls(nodeFrom, methodName)
+  // )
+  // or
   // list.append, set.add
   exists(DataFlow::MethodCallNode call, DataFlow::Node obj |
     call.calls(obj, ["append", "add"]) and
     obj = nodeTo.(DataFlow::PostUpdateNode).getPreUpdateNode() and
     call.getArg(0) = nodeFrom
   )
+  or
+  // Although flow through collections is modeled precisely using stores/reads, we still
+  // allow flow out of a _tainted_ collection. This is needed in order to support taint-
+  // tracking configurations where the source is a collection.
+  exists(DataFlow::Content c | DataFlowPrivate::readStep(nodeFrom, c, nodeTo) |
+    c instanceof DataFlow::ListElementContent
+    or
+    c instanceof DataFlow::SetElementContent
+    or
+    c instanceof DataFlow::DictionaryElementContent
+    or
+    c instanceof DataFlow::DictionaryElementAnyContent
+  )
 }
 
 /**