Check for wrapper classes

python/ql/src/Resources/FileNotAlwaysClosedQuery.qll

@@ -1,22 +1,36 @@
 /** Definitions for reasoning about whether files are closed. */
 
 import python
-//import semmle.python.dataflow.DataFlow
+import semmle.python.dataflow.new.internal.DataFlowDispatch
 import semmle.python.ApiGraphs
 
 abstract class FileOpen extends DataFlow::CfgNode { }
 
 class FileOpenCall extends FileOpen {
-  FileOpenCall() { this = API::builtin("open").getACall() }
+  FileOpenCall() { this = [API::builtin("open").getACall()] }
+}
+
+class FileWrapperClassCall extends FileOpen, DataFlow::CallCfgNode {
+  FileOpen wrapped;
+
+  FileWrapperClassCall() {
+    wrapped = this.getArg(_).getALocalSource() and
+    this.getFunction() = classTracker(_)
+  }
+
+  FileOpen getWrapped() { result = wrapped }
 }
 
-// todo: type tracking to find wrapping funcs
 abstract class FileClose extends DataFlow::CfgNode { }
 
 class FileCloseCall extends FileClose {
   FileCloseCall() { exists(DataFlow::MethodCallNode mc | mc.calls(this, "close")) }
 }
 
+class OsCloseCall extends FileClose {
+  OsCloseCall() { this = API::moduleImport("os").getMember("close").getACall().getArg(0) }
+}
+
 class WithStatement extends FileClose {
   WithStatement() { exists(With w | this.asExpr() = w.getContextExpr()) }
 }
@@ -34,6 +48,6 @@ predicate fileIsStoredInField(FileOpen fo) {
 predicate fileNotAlwaysClosed(FileOpen fo) {
   not fileIsClosed(fo) and
   not fileIsReturned(fo) and
-  not fileIsStoredInField(fo)
-  // TODO: exception cases
+  not fileIsStoredInField(fo) and
+  not exists(FileWrapperClassCall fwc | fo = fwc.getWrapped())
 }