Python: Add importModule and importMember DataFlow helpers

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

@@ -16,3 +16,52 @@ predicate localFlowStep(Node nodeFrom, Node nodeTo) { simpleLocalFlowStep(nodeFr
  * (intra-procedural) steps.
  */
 predicate localFlow(Node source, Node sink) { localFlowStep*(source, sink) }
+
+/**
+ * Gets an EssaNode that holds the module imported by `name`.
+ * Note that for the statement `import pkg.mod`, the new variable introduced is `pkg` that is a
+ * reference to the module `pkg`.
+ *
+ * This predicate handles (with optional `... as <new-name>`):
+ * 1. `import <name>`
+ * 2. `from <package> import <module>` when `<name> = <package> + "." + <module>`
+ * 3. `from <module> import <member>` when `<name> = <module> + "." + <member>`
+ *
+ * Note:
+ * While it is technically possible that `import mypkg.foo` and `from mypkg import foo` can give different values,
+ * it's highly unlikely that this will be a problem in production level code.
+ *   Example: If `mypkg/__init__.py` contains `foo = 42`, then `from mypkg import foo` will not import the module
+ *   `mypkg/foo.py` but the variable `foo` containing `42` -- however, `import mypkg.foo` will always cause `mypkg.foo`
+ *   to refer to the module.
+ *
+ * Also see `DataFlow::importMember`
+ */
+EssaNode importModule(string name) {
+  exists(Variable var, Import imp, Alias alias |
+    alias = imp.getAName() and
+    alias.getAsname() = var.getAStore() and
+    (
+      name = alias.getValue().(ImportMember).getImportedModuleName()
+      or
+      name = alias.getValue().(ImportExpr).getImportedModuleName()
+    ) and
+    result.getVar().(AssignmentDefinition).getSourceVariable() = var
+  )
+}
+
+/**
+ * Gets a EssaNode that holds the value imported by using fully qualified name in
+ *`from <moduleName> import <memberName>`.
+ *
+ * Also see `DataFlow::importModule`.
+ */
+EssaNode importMember(string moduleName, string memberName) {
+  exists(Variable var, Import imp, Alias alias, ImportMember member |
+    alias = imp.getAName() and
+    member = alias.getValue() and
+    moduleName = member.getModule().(ImportExpr).getImportedModuleName() and
+    memberName = member.getName() and
+    alias.getAsname() = var.getAStore() and
+    result.getVar().(AssignmentDefinition).getSourceVariable() = var
+  )
+}

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

@@ -0,0 +1,20 @@
+importModule
+| test1.py:1:8:1:12 | GSSA Variable mypkg | 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:2:8:2:16 | GSSA Variable mypkg | mypkg |
+| test4.py:1:21:1:24 | GSSA Variable _foo | mypkg.foo |
+| test4.py:2:21:2:24 | GSSA Variable _bar | mypkg.bar |
+| test5.py:1:8:1:12 | GSSA Variable mypkg | mypkg |
+| test5.py:9:26:9:29 | GSSA Variable _bar | mypkg.bar |
+| test6.py:1:8:1:12 | GSSA Variable mypkg | mypkg |
+| test6.py:5:8:5:16 | GSSA Variable mypkg | mypkg |
+| test7.py:1:19:1:21 | GSSA Variable foo | mypkg.foo |
+| test7.py:5:8:5:16 | GSSA Variable mypkg | mypkg |
+| test7.py:9:19:9:21 | GSSA Variable foo | mypkg.foo |
+importMember
+| test2.py:1:19:1:21 | GSSA Variable foo | mypkg | foo |
+| test2.py:1:24:1:26 | GSSA Variable bar | mypkg | bar |
+| test5.py:9:26:9:29 | GSSA Variable _bar | mypkg | bar |
+| test7.py:1:19:1:21 | GSSA Variable foo | mypkg | foo |
+| test7.py:9:19:9:21 | GSSA Variable foo | mypkg | foo |

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

@@ -0,0 +1,8 @@
+import python
+import experimental.dataflow.DataFlow
+
+query predicate importModule(DataFlow::Node res, string name) { res = DataFlow::importModule(name) }
+
+query predicate importMember(DataFlow::Node res, string moduleName, string memberName) {
+  res = DataFlow::importMember(moduleName, memberName)
+}

python/ql/test/experimental/dataflow/import-helper/README.md

@@ -0,0 +1 @@
+Small tests that explore difference between `import mypkg.foo` and `from mypkg import foo`.

python/ql/test/experimental/dataflow/import-helper/mypkg/__init__.py

@@ -0,0 +1 @@
+foo = 42

python/ql/test/experimental/dataflow/import-helper/mypkg/bar.py

@@ -0,0 +1 @@
+pass

python/ql/test/experimental/dataflow/import-helper/mypkg/foo.py

@@ -0,0 +1 @@
+pass

python/ql/test/experimental/dataflow/import-helper/test1.py

@@ -0,0 +1,6 @@
+import mypkg
+print(mypkg.foo)  # 42
+try:
+    print(mypkg.bar)
+except AttributeError as e:
+    print(e)  # module 'mypkg' has no attribute 'bar'

python/ql/test/experimental/dataflow/import-helper/test2.py

@@ -0,0 +1,3 @@
+from mypkg import foo, bar
+print(foo)
+print(bar)

python/ql/test/experimental/dataflow/import-helper/test3.py

@@ -0,0 +1,4 @@
+import mypkg.foo
+import mypkg.bar
+print(mypkg.foo)  # <module 'mypkg.foo' ...
+print(mypkg.bar)  # <module 'mypkg.bar' ...

python/ql/test/experimental/dataflow/import-helper/test4.py

@@ -0,0 +1,4 @@
+import mypkg.foo as _foo
+import mypkg.bar as _bar
+print(_foo)  # <module 'mypkg.bar' ...
+print(_bar)  # <module 'mypkg.bar' ...

python/ql/test/experimental/dataflow/import-helper/test5.py

@@ -0,0 +1,10 @@
+import mypkg
+
+print(mypkg.foo)  # 42
+try:
+    print(mypkg.bar)
+except AttributeError as e:
+    print(e)  # module 'mypkg' has no attribute 'bar'
+
+from mypkg import bar as _bar
+print(mypkg.bar)  # <module 'mypkg.bar' ...

python/ql/test/experimental/dataflow/import-helper/test6.py

@@ -0,0 +1,6 @@
+import mypkg
+
+print(mypkg.foo) # 42
+
+import mypkg.foo
+print(mypkg.foo) # <module 'mypkg.foo' ...

python/ql/test/experimental/dataflow/import-helper/test7.py

@@ -0,0 +1,10 @@
+from mypkg import foo
+
+print(foo)  # 42
+
+import mypkg.foo
+print(foo)  # 42
+print(mypkg.foo)  # <module 'mypkg.bar' ...
+
+from mypkg import foo
+print(foo)  # <module 'mypkg.bar' ...