Python: Support named imports as attribute reads

Required a small change in `DataFlow::importModule` to get the desired
behaviour (cf. the type trackers defined in `moduleattr.ql`, but this
should be harmless. The node that is added doesn't have any flow
anywhere.

Author: tausbn

python/ql/src/experimental/dataflow/internal/Attributes.qll

@@ -218,3 +218,31 @@ private class GetAttrCallAsAttrRead extends AttrRead, CfgNode {
     result = this.getAttributeNameExpr().asExpr().(StrConst).getText()
   }
 }
+
+/**
+ * A convenience class for embedding `ImportMemberNode` into `DataFlowCfgNode`, as the former is not
+ * obviously a subtype of the latter.
+ */
+private class DataFlowImportMemberNode extends ImportMemberNode, DataFlowCfgNode { }
+
+/**
+ * Represents a named import as an attribute read. That is,
+ * ```python
+ * from module import attr as attr_ref
+ * ```
+ * is treated as if it is a read of the attribute `module.attr`, even if `module` is not imported directly.
+ */
+private class ModuleAttributeImportAsAttrRead extends AttrRead, CfgNode {
+  override DataFlowImportMemberNode node;
+
+  override Node getObject() { result.asCfgNode() = node.getModule(_) }
+
+  override ExprNode getAttributeNameExpr() {
+    // The name of an imported attribute doesn't exist as a `Node` in the control flow graph, as it
+    // can only ever be an identifier, and is therefore represented directly as a string.
+    // Use `getAttributeName` to access the name of the attribute.
+    none()
+  }
+
+  override string getAttributeName() { exists(node.getModule(result)) }
+}

python/ql/src/experimental/dataflow/internal/DataFlowUtil.qll

@@ -36,7 +36,7 @@ predicate localFlow(Node source, Node sink) { localFlowStep*(source, sink) }
  *
  * Also see `DataFlow::importMember`
  */
-EssaNode importModule(string name) {
+Node importModule(string name) {
   exists(Variable var, Import imp, Alias alias |
     alias = imp.getAName() and
     alias.getAsname() = var.getAStore() and
@@ -45,8 +45,14 @@ EssaNode importModule(string name) {
       or
       name = alias.getValue().(ImportExpr).getImportedModuleName()
     ) and
-    result.getVar().(AssignmentDefinition).getSourceVariable() = var
+    result.(EssaNode).getVar().(AssignmentDefinition).getSourceVariable() = var
   )
+  or
+  // In `from module import attr`, we want to consider `module` to be an expression that refers to a
+  // module of that name, as this allows us to refer to attributes of this module, even if it's
+  // never imported directly. Note that there crucially isn't any _flow_ from `module` to references
+  // to that same identifier.
+  result.asCfgNode().getNode() = any(ImportExpr i | i.getAnImportedModuleName() = name)
 }
 
 /**

python/ql/test/experimental/dataflow/import-helper/ImportHelper.expected

@@ -1,16 +1,36 @@
 importModule
+| test1.py:1:8:1:12 | ControlFlowNode for ImportExpr | mypkg |
 | test1.py:1:8:1:12 | GSSA Variable mypkg | mypkg |
+| test2.py:1:6:1:10 | ControlFlowNode for ImportExpr | mypkg |
+| test2.py:1:6:1:10 | ControlFlowNode for ImportExpr | mypkg |
 | test2.py:1:19:1:21 | GSSA Variable foo | mypkg.foo |
 | test2.py:1:24:1:26 | GSSA Variable bar | mypkg.bar |
+| test3.py:1:8:1:16 | ControlFlowNode for ImportExpr | mypkg |
+| test3.py:1:8:1:16 | ControlFlowNode for ImportExpr | mypkg.foo |
+| test3.py:2:8:2:16 | ControlFlowNode for ImportExpr | mypkg |
+| test3.py:2:8:2:16 | ControlFlowNode for ImportExpr | mypkg.bar |
 | test3.py:2:8:2:16 | GSSA Variable mypkg | mypkg |
+| test4.py:1:8:1:16 | ControlFlowNode for ImportExpr | mypkg |
+| test4.py:1:8:1:16 | ControlFlowNode for ImportExpr | mypkg.foo |
 | test4.py:1:21:1:24 | GSSA Variable _foo | mypkg.foo |
+| test4.py:2:8:2:16 | ControlFlowNode for ImportExpr | mypkg |
+| test4.py:2:8:2:16 | ControlFlowNode for ImportExpr | mypkg.bar |
 | test4.py:2:21:2:24 | GSSA Variable _bar | mypkg.bar |
+| test5.py:1:8:1:12 | ControlFlowNode for ImportExpr | mypkg |
 | test5.py:1:8:1:12 | GSSA Variable mypkg | mypkg |
+| test5.py:9:6:9:10 | ControlFlowNode for ImportExpr | mypkg |
 | test5.py:9:26:9:29 | GSSA Variable _bar | mypkg.bar |
+| test6.py:1:8:1:12 | ControlFlowNode for ImportExpr | mypkg |
 | test6.py:1:8:1:12 | GSSA Variable mypkg | mypkg |
+| test6.py:5:8:5:16 | ControlFlowNode for ImportExpr | mypkg |
+| test6.py:5:8:5:16 | ControlFlowNode for ImportExpr | mypkg.foo |
 | test6.py:5:8:5:16 | GSSA Variable mypkg | mypkg |
+| test7.py:1:6:1:10 | ControlFlowNode for ImportExpr | mypkg |
 | test7.py:1:19:1:21 | GSSA Variable foo | mypkg.foo |
+| test7.py:5:8:5:16 | ControlFlowNode for ImportExpr | mypkg |
+| test7.py:5:8:5:16 | ControlFlowNode for ImportExpr | mypkg.foo |
 | test7.py:5:8:5:16 | GSSA Variable mypkg | mypkg |
+| test7.py:9:6:9:10 | ControlFlowNode for ImportExpr | mypkg |
 | test7.py:9:19:9:21 | GSSA Variable foo | mypkg.foo |
 importMember
 | test2.py:1:19:1:21 | GSSA Variable foo | mypkg | foo |

python/ql/test/experimental/dataflow/typetracking/import_as_attr.py

@@ -0,0 +1,9 @@
+from module import attr as attr_ref
+
+x = attr_ref
+
+def fun():
+    y = attr_ref
+
+# The following should _not_ be a reference to the above module, since we don't actually import it.
+z = module

python/ql/test/experimental/dataflow/typetracking/moduleattr.expected

@@ -0,0 +1,10 @@
+module_tracker
+| import_as_attr.py:1:6:1:11 | ControlFlowNode for ImportExpr |
+module_attr_tracker
+| import_as_attr.py:0:0:0:0 | ModuleVariableNode for Global Variable attr_ref in Module import_as_attr |
+| import_as_attr.py:1:20:1:35 | ControlFlowNode for ImportMember |
+| import_as_attr.py:1:28:1:35 | GSSA Variable attr_ref |
+| import_as_attr.py:3:1:3:1 | GSSA Variable x |
+| import_as_attr.py:3:5:3:12 | ControlFlowNode for attr_ref |
+| import_as_attr.py:6:5:6:5 | SSA variable y |
+| import_as_attr.py:6:9:6:16 | ControlFlowNode for attr_ref |

python/ql/test/experimental/dataflow/typetracking/moduleattr.ql

@@ -0,0 +1,23 @@
+import python
+import experimental.dataflow.DataFlow
+import experimental.dataflow.TypeTracker
+
+DataFlow::Node module_tracker(TypeTracker t) {
+  t.start() and
+  result = DataFlow::importModule("module")
+  or
+  exists(TypeTracker t2 | result = module_tracker(t2).track(t2, t))
+}
+
+query DataFlow::Node module_tracker() { result = module_tracker(DataFlow::TypeTracker::end()) }
+
+DataFlow::Node module_attr_tracker(TypeTracker t) {
+  t.startInAttr("attr") and
+  result = module_tracker()
+  or
+  exists(TypeTracker t2 | result = module_attr_tracker(t2).track(t2, t))
+}
+
+query DataFlow::Node module_attr_tracker() {
+  result = module_attr_tracker(DataFlow::TypeTracker::end())
+}