Python: Move import logic into its own module

tausbn · tausbn · commit 366c5743084a · 2022-09-08T14:52:08.000Z
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll
@@ -2,6 +2,7 @@ private import python
 private import DataFlowPublic
 private import semmle.python.essa.SsaCompute
 private import semmle.python.dataflow.new.internal.ImportStar
+private import semmle.python.dataflow.new.internal.ImportResolution
 // Since we allow extra data-flow steps from modeled frameworks, we import these
 // up-front, to ensure these are included. This provides a more seamless experience from
 // a user point of view, since they don't need to know they need to import a specific
@@ -419,9 +420,9 @@ predicate jumpStepSharedWithTypeTracker(Node nodeFrom, Node nodeTo) {
   runtimeJumpStep(nodeFrom, nodeTo)
   or
   // Read of module attribute:
-  exists(AttrRead r, ModuleValue mv |
-    r.getObject().asCfgNode().pointsTo(mv) and
-    module_export(mv.getScope(), r.getAttributeName(), nodeFrom) and
+  exists(AttrRead r |
+    ImportResolution::module_export(ImportResolution::getModule(r.getObject()),
+      r.getAttributeName(), nodeFrom) and
     nodeTo = r
   )
   or
@@ -445,22 +446,6 @@ predicate jumpStepNotSharedWithTypeTracker(Node nodeFrom, Node nodeTo) {
   any(Orm::AdditionalOrmSteps es).jumpStep(nodeFrom, nodeTo)
 }
 
-/**
- * Holds if the module `m` defines a name `name` by assigning `defn` to it. This is an
- * overapproximation, as `name` may not in fact be exported (e.g. by defining an `__all__` that does
- * not include `name`).
- */
-private predicate module_export(Module m, string name, CfgNode defn) {
-  exists(EssaVariable v |
-    v.getName() = name and
-    v.getAUse() = ImportStar::getStarImported*(m).getANormalExit()
-  |
-    defn.getNode() = v.getDefinition().(AssignmentDefinition).getValue()
-    or
-    defn.getNode() = v.getDefinition().(ArgumentRefinement).getArgument()
-  )
-}
-
 //--------
 // Field flow
 //--------
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/ImportResolution.qll b/python/ql/lib/semmle/python/dataflow/new/internal/ImportResolution.qll
@@ -0,0 +1,29 @@
+private import python
+private import semmle.python.dataflow.new.DataFlow
+private import semmle.python.dataflow.new.internal.ImportStar
+private import semmle.python.dataflow.new.TypeTracker
+
+module ImportResolution {
+  /**
+   * Holds if the module `m` defines a name `name` by assigning `defn` to it. This is an
+   * overapproximation, as `name` may not in fact be exported (e.g. by defining an `__all__` that does
+   * not include `name`).
+   */
+  predicate module_export(Module m, string name, DataFlow::CfgNode defn) {
+    exists(EssaVariable v |
+      v.getName() = name and
+      v.getAUse() = ImportStar::getStarImported*(m).getANormalExit()
+    |
+      defn.getNode() = v.getDefinition().(AssignmentDefinition).getValue()
+      or
+      defn.getNode() = v.getDefinition().(ArgumentRefinement).getArgument()
+    )
+  }
+
+  Module getModule(DataFlow::CfgNode node) {
+    exists(ModuleValue mv |
+      node.getNode().pointsTo(mv) and
+      result = mv.getScope()
+    )
+  }
+}
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/ImportStar.qll b/python/ql/lib/semmle/python/dataflow/new/internal/ImportStar.qll
@@ -2,6 +2,8 @@
 
 private import python
 private import semmle.python.dataflow.new.internal.Builtins
+private import semmle.python.dataflow.new.internal.ImportResolution
+private import semmle.python.dataflow.new.DataFlow
 
 cached
 module ImportStar {
@@ -71,8 +73,10 @@ module ImportStar {
    */
   cached
   Module getStarImported(Module m) {
-    exists(ImportStar i |
-      i.getScope() = m and result = i.getModule().pointsTo().(ModuleValue).getScope()
+    exists(ImportStar i, DataFlow::CfgNode imported_module |
+      imported_module.getNode().getNode() = i.getModule() and
+      i.getScope() = m and
+      result = ImportResolution::getModule(imported_module)
     )
   }
 
