Python: Recognize shebangs in module usage detection

RasmusWL · RasmusWL · commit 82b29b56988e · 2020-02-19T14:12:22.000+01:00
diff --git a/python/ql/src/semmle/python/objects/ObjectAPI.qll b/python/ql/src/semmle/python/objects/ObjectAPI.qll
@@ -210,7 +210,12 @@ class ModuleValue extends Value {
                 i.getTest().(Compare).compares(name, op, main) and
                 name.getId() = "__name__" and main.getText() = "__main__"
             )
-            // TODO: Add she-bang handling
+            or
+            exists(Comment c |
+                c.getLocation().getFile() = this.getPath() and
+                c.getLocation().getStartLine() = 1 and
+                c.getText().regexpMatch("^#!/.*python(2|3)?[ \\\\t]*$")
+            )
         )
     }
 }
diff --git a/python/ql/test/library-tests/modules/usage/ModuleUsage.expected b/python/ql/test/library-tests/modules/usage/ModuleUsage.expected
@@ -2,5 +2,5 @@
 | 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 |
-| myscript.py:0:0:0:0 | Script myscript | <UNKNOWN> |
+| myscript.py:0:0:0:0 | Script myscript | isUsedAsScript |
 | script:0:0:0:0 | Script script | isUsedAsScript |

Original file line number	Diff line number	Diff line change
`@@ -210,7 +210,12 @@ class ModuleValue extends Value {`
`210`	`210`	`i.getTest().(Compare).compares(name, op, main) and`
`211`	`211`	`name.getId() = "__name__" and main.getText() = "__main__"`
`212`	`212`	`)`
`213`		`- // TODO: Add she-bang handling`
	`213`	`+ or`
	`214`	`+ exists(Comment c \|`
	`215`	`+ c.getLocation().getFile() = this.getPath() and`
	`216`	`+ c.getLocation().getStartLine() = 1 and`
	`217`	`+ c.getText().regexpMatch("^#!/.python(2\|3)?[ \\\\t]$")`
	`218`	`+ )`
`214`	`219`	`)`
`215`	`220`	`}`
`216`	`221`	`}`