Python: Allow absolute imports in directories with scripts

Fixes the import logic to account for absolute imports.

We do this by classifying which files and folders may serve as the
entry point for execution, based on a few simple heuristics. If the
file `module.py` is in the same folder as a file `main.py` that may be
executed directly, then we allow `module` to be a valid name for
`module.py` so that `import module` will work as expected.

Author: tausbn

python/ql/src/semmle/python/Files.qll

@@ -72,6 +72,33 @@ class File extends Container {
    * are specified to be extracted.
    */
   string getContents() { file_contents(this, result) }
+
+  /** Holds if this file is likely to get executed directly, and thus act as an entry point for execution. */
+  predicate maybeExecutedDirectly() {
+    // Only consider files in the source code, and not things like the standard library
+    exists(this.getRelativePath()) and
+    (
+      // The file doesn't have the extension `.py` but still contains Python statements
+      not this.getExtension() = "py" and
+      exists(Stmt s | s.getLocation().getFile() = this)
+      or
+      // The file contains the usual `if __name__ == '__main__':` construction
+      exists(If i, Name name, StrConst main, Cmpop op |
+        i.getScope().(Module).getFile() = this and
+        op instanceof Eq and
+        i.getTest().(Compare).compares(name, op, main) and
+        name.getId() = "__name__" and
+        main.getText() = "__main__"
+      )
+      or
+      // The file contains a `#!` line referencing the python interpreter
+      exists(Comment c |
+        c.getLocation().getFile() = this and
+        c.getLocation().getStartLine() = 1 and
+        c.getText().regexpMatch("^#! */.*python(2|3)?[ \\\\t]*$")
+      )
+    )
+  }
 }
 
 private predicate occupied_line(File f, int n) {
@@ -121,6 +148,9 @@ class Folder extends Container {
     this.getBaseName().regexpMatch("[^\\d\\W]\\w*") and
     result = this.getParent().getImportRoot(n)
   }
+
+  /** Holds if execution may start in a file in this directory. */
+  predicate mayContainEntryPoint() { any(File f | f.getParent() = this).maybeExecutedDirectly() }
 }
 
 /**

python/ql/src/semmle/python/Module.qll

@@ -204,8 +204,13 @@ private string moduleNameFromBase(Container file) {
 string moduleNameFromFile(Container file) {
   exists(string basename |
     basename = moduleNameFromBase(file) and
-    legalShortName(basename) and
+    legalShortName(basename)
+  |
     result = moduleNameFromFile(file.getParent()) + "." + basename
+    or
+    // If execution can start in the folder containing this module, then we will assume `file` can
+    // be imported as an absolute import, and hence return `basename` as a possible name.
+    file.getParent().(Folder).mayContainEntryPoint() and result = basename
   )
   or
   isPotentialSourcePackage(file) and

python/ql/test/3/library-tests/modules/entry_point/modules.expected

@@ -1,3 +1,5 @@
+| main | code/main.py:0:0:0:0 | Script main |
+| module | code/module.py:0:0:0:0 | Module module |
 | package | code/package:0:0:0:0 | Package package |
 | package.__init__ | code/package/__init__.py:0:0:0:0 | Module package.__init__ |
 | package.package_main | code/package/package_main.py:0:0:0:0 | Module package.package_main |