Python: Allow import resolution with recursive phi/refine steps

Author: RasmusWL

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

@@ -64,6 +64,19 @@ private import semmle.python.dataflow.new.internal.DataFlowPrivate
  *   `bar` subpackage.
  */
 module ImportResolution {
+  /**
+   * Holds if there is an ESSA step from `defFrom` to `defTo`, which should be allowed
+   * for import resolution.
+   */
+  private predicate allowedEssaImportStep(EssaDefinition defFrom, EssaDefinition defTo) {
+    // to handle definitions guarded by if-then-else
+    defFrom = defTo.(PhiFunction).getAnInput()
+    or
+    // refined variable
+    // example: https://github.com/nvbn/thefuck/blob/ceeaeab94b5df5a4fe9d94d61e4f6b0bbea96378/thefuck/utils.py#L25-L45
+    defFrom = defTo.(EssaNodeRefinement).getInput().getDefinition()
+  }
+
   /**
    * 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
@@ -74,15 +87,7 @@ module ImportResolution {
     exists(EssaVariable v, EssaDefinition essaDef |
       v.getName() = name and
       v.getAUse() = ImportStar::getStarImported*(m).getANormalExit() and
-      (
-        essaDef = v.getDefinition()
-        or
-        // to handle definitions guarded by if-then-else
-        essaDef = v.getDefinition().(PhiFunction).getAnInput()
-        or
-        // refined variable
-        essaDef = v.getDefinition().(EssaNodeRefinement).getInput().getDefinition()
-      )
+      allowedEssaImportStep*(essaDef, v.getDefinition())
     |
       defn.getNode() = essaDef.(AssignmentDefinition).getValue()
       or

python/ql/test/experimental/import-resolution/main.py

@@ -95,10 +95,10 @@ def local_import():
 
 # check that refined definitions are handled correctly
 import refined # $ imports=refined as=refined
-check("refined.SOURCE", refined.SOURCE, refined.SOURCE, globals()) #$ MISSING: prints=SOURCE
+check("refined.SOURCE", refined.SOURCE, refined.SOURCE, globals()) #$ prints=SOURCE
 
 import if_then_else_refined # $ imports=if_then_else_refined as=if_then_else_refined
-check("if_then_else_refined.src", if_then_else_refined.src, if_then_else_refined.src, globals()) #$ MISSING: prints=SOURCE
+check("if_then_else_refined.src", if_then_else_refined.src, if_then_else_refined.src, globals()) #$ prints=SOURCE
 
 exit(__file__)