Python: Add ModuleValue.{isUsedAsModule, isUsedAsScript}

RasmusWL · RasmusWL · commit 3e7e9636eaa8 · 2020-02-19T14:12:22.000+01:00
and a few test cases
diff --git a/python/ql/src/Statements/TopLevelPrint.ql b/python/ql/src/Statements/TopLevelPrint.ql
@@ -35,6 +35,6 @@ from Stmt p
 where
     is_print_stmt(p) and
     // TODO: Need to discuss how we would like to handle ModuleObject.getKind in the glorious future
-    exists(ModuleObject m | m.getModule() = p.getScope() and m.getKind() = "module") and
+    exists(ModuleValue m | m.getScope() = p.getScope() and m.isUsedAsModule()) and
     not exists(If i | main_eq_name(i) and i.getASubStatement().getASubStatement*() = p)
 select p, "Print statement may execute during import."
diff --git a/python/ql/src/semmle/python/objects/ObjectAPI.qll b/python/ql/src/semmle/python/objects/ObjectAPI.qll
@@ -187,6 +187,32 @@ class ModuleValue extends Value {
         result.importedAs(this.getScope().getAnImportedModuleName())
     }
 
+    /** When used as a normal module (for example, imported and used by other modules) */
+    predicate isUsedAsModule() {
+        this.isBuiltin()
+        or
+        this.isPackage()
+        or
+        exists(ImportingStmt i | this.importedAs(i.getAnImportedModuleName()))
+        or
+        this.getPath().getBaseName() = "__init__.py"
+    }
+
+    /** When used (exclusively) as a script (will not include normal modules that can also be run as a script) */
+    predicate isUsedAsScript() {
+        not isUsedAsModule() and
+        (
+            not this.getPath().getExtension() = "py"
+            or
+            exists(If i, Name name, StrConst main, Cmpop op |
+                i.getScope() = this.getScope() and
+                op instanceof Eq and
+                i.getTest().(Compare).compares(name, op, main) and
+                name.getId() = "__name__" and main.getText() = "__main__"
+            )
+            // TODO: Add she-bang handling
+        )
+    }
 }
 
 module Module {
diff --git a/python/ql/test/library-tests/modules/usage/ModuleUsage.expected b/python/ql/test/library-tests/modules/usage/ModuleUsage.expected
@@ -0,0 +1,5 @@
+| file://:0:0:0:0 | Module builtins | isUsedAsModule |
+| file://:0:0:0:0 | Module sys | isUsedAsModule |
+| imported.py:0:0:0:0 | Module imported | isUsedAsModule |
+| main.py:0:0:0:0 | Module main | isUsedAsScript |
+| mybin:0:0:0:0 | Script script | isUsedAsScript |
diff --git a/python/ql/test/library-tests/modules/usage/ModuleUsage.ql b/python/ql/test/library-tests/modules/usage/ModuleUsage.ql
@@ -0,0 +1,12 @@
+import python
+
+from ModuleValue mv, string usage
+where
+    mv.isUsedAsModule() and usage = "isUsedAsModule"
+    or
+    mv.isUsedAsScript() and usage = "isUsedAsScript"
+    or
+    not mv.isUsedAsModule() and
+    not mv.isUsedAsScript() and
+    usage = "<UNKNOWN>"
+select mv, usage
diff --git a/python/ql/test/library-tests/modules/usage/imported.py b/python/ql/test/library-tests/modules/usage/imported.py
@@ -0,0 +1,6 @@
+def func():
+    pass
+
+if __name__ == "__main__":
+    print("I could have done something interesting...")
+    print("but I didn't")
diff --git a/python/ql/test/library-tests/modules/usage/main.py b/python/ql/test/library-tests/modules/usage/main.py
@@ -0,0 +1,5 @@
+import imported
+
+if __name__ == "__main__":
+    imported.func()
+    print('Done')
diff --git a/python/ql/test/library-tests/modules/usage/options b/python/ql/test/library-tests/modules/usage/options
@@ -0,0 +1 @@
+semmle-extractor-options: -F script
diff --git a/python/ql/test/library-tests/modules/usage/script b/python/ql/test/library-tests/modules/usage/script
@@ -0,0 +1,3 @@
+#!/usr/bin/env python
+
+print('Under construction :)')

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+#!/usr/bin/env python`
	`2`	`+`
	`3`	`+print('Under construction :)')`