Python: Rewrite most of SQLAlchemy modeling

Author: RasmusWL

python/ql/src/experimental/semmle/python/frameworks/SqlAlchemy.qll

@@ -1,6 +1,8 @@
 /**
- * Provides classes modeling security-relevant aspects of the 'SqlAlchemy' package.
- * See https://pypi.org/project/SQLAlchemy/.
+ * Provides classes modeling security-relevant aspects of the `SQLAlchemy` PyPI package.
+ * See
+ *  - https://pypi.org/project/SQLAlchemy/
+ *  - https://docs.sqlalchemy.org/en/14/index.html
  */
 
 private import python
@@ -10,93 +12,209 @@ private import semmle.python.ApiGraphs
 private import semmle.python.Concepts
 private import experimental.semmle.python.Concepts
 
+/**
+ * Provides models for the `SQLAlchemy` PyPI package.
+ * See
+ *  - https://pypi.org/project/SQLAlchemy/
+ *  - https://docs.sqlalchemy.org/en/14/index.html
+ */
 private module SqlAlchemy {
   /**
-   * Returns an instantization of a SqlAlchemy Session object.
-   * See https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.Session and
-   * https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.sessionmaker
+   * Provides models for the `sqlalchemy.engine.Engine` and `sqlalchemy.future.Engine` classes.
+   *
+   * These are so similar that we model both in the same way.
+   *
+   * See
+   *  - https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.Engine
+   *  - https://docs.sqlalchemy.org/en/14/core/future.html#sqlalchemy.future.Engine
    */
-  private API::Node getSqlAlchemySessionInstance() {
-    result = API::moduleImport("sqlalchemy.orm").getMember("Session").getReturn() or
-    result = API::moduleImport("sqlalchemy.orm").getMember("sessionmaker").getReturn().getReturn()
-  }
+  module Engine {
+    /** Gets a reference to a SQLAlchemy Engine class. */
+    private API::Node classRef() {
+      result = API::moduleImport("sqlalchemy").getMember("engine").getMember("Engine")
+      or
+      result = API::moduleImport("sqlalchemy").getMember("future").getMember("Engine")
+    }
 
-  /**
-   * Returns an instantization of a SqlAlchemy Engine object.
-   * See https://docs.sqlalchemy.org/en/14/core/engines.html#sqlalchemy.create_engine
-   */
-  private API::Node getSqlAlchemyEngineInstance() {
-    result = API::moduleImport("sqlalchemy").getMember("create_engine").getReturn()
-  }
+    /**
+     * A source of instances of a SQLAlchemy Engine, extend this class to model new instances.
+     *
+     * This can include instantiations of the class, return values from function
+     * calls, or a special parameter that will be set when functions are called by an external
+     * library.
+     *
+     * Use the predicate `Engine::instance()` to get references to instances of a SQLAlchemy Engine.
+     */
+    abstract class InstanceSource extends DataFlow::LocalSourceNode { }
 
-  /**
-   * Returns an instantization of a SqlAlchemy Query object.
-   * See https://docs.sqlalchemy.org/en/14/orm/query.html?highlight=query#sqlalchemy.orm.Query
-   */
-  private API::Node getSqlAlchemyQueryInstance() {
-    result = getSqlAlchemySessionInstance().getMember("query").getReturn()
+    private class EngineConstruction extends InstanceSource, DataFlow::CallCfgNode {
+      EngineConstruction() {
+        this = classRef().getACall()
+        or
+        this = API::moduleImport("sqlalchemy").getMember("create_engine").getACall()
+        or
+        this =
+          API::moduleImport("sqlalchemy").getMember("future").getMember("create_engine").getACall()
+        or
+        this.(DataFlow::MethodCallNode).calls(instance(), "execution_options")
+      }
+    }
+
+    /** Gets a reference to an instance of a SQLAlchemy Engine. */
+    private DataFlow::TypeTrackingNode instance(DataFlow::TypeTracker t) {
+      t.start() and
+      result instanceof InstanceSource
+      or
+      exists(DataFlow::TypeTracker t2 | result = instance(t2).track(t2, t))
+    }
+
+    /** Gets a reference to an instance of a SQLAlchemy Engine. */
+    DataFlow::Node instance() { instance(DataFlow::TypeTracker::end()).flowsTo(result) }
   }
 
   /**
-   * A call to `execute` meant to execute an SQL expression
-   * See the following links:
-   *   - https://docs.sqlalchemy.org/en/14/core/connections.html?highlight=execute#sqlalchemy.engine.Connection.execute
-   *   - https://docs.sqlalchemy.org/en/14/core/connections.html?highlight=execute#sqlalchemy.engine.Engine.execute
-   *   - https://docs.sqlalchemy.org/en/14/orm/session_api.html?highlight=execute#sqlalchemy.orm.Session.execute
+   * Provides models for the `sqlalchemy.engine.base.Connection` and `sqlalchemy.future.Connection` classes.
+   *
+   * These are so similar that we model both in the same way.
+   *
+   * See
+   * - https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.Connection
+   * - https://docs.sqlalchemy.org/en/14/core/future.html#sqlalchemy.future.Connection
    */
-  private class SqlAlchemyExecuteCall extends DataFlow::CallCfgNode, SqlExecution::Range {
-    SqlAlchemyExecuteCall() {
-      // new way
-      this = getSqlAlchemySessionInstance().getMember("execute").getACall() or
-      this =
-        getSqlAlchemyEngineInstance()
-            .getMember(["connect", "begin"])
-            .getReturn()
-            .getMember("execute")
-            .getACall()
+  module Connection {
+    /** Gets a reference to a SQLAlchemy Connection class. */
+    private API::Node classRef() {
+      result =
+        API::moduleImport("sqlalchemy")
+            .getMember("engine")
+            .getMember("base")
+            .getMember("Connection")
+      or
+      result = API::moduleImport("sqlalchemy").getMember("future").getMember("Connection")
     }
 
-    override DataFlow::Node getSql() { result = this.getArg(0) }
+    /**
+     * A source of instances of a SQLAlchemy Connection, extend this class to model new instances.
+     *
+     * This can include instantiations of the class, return values from function
+     * calls, or a special parameter that will be set when functions are called by an external
+     * library.
+     *
+     * Use the predicate `Connection::instance()` to get references to instances of a SQLAlchemy Connection.
+     */
+    abstract class InstanceSource extends DataFlow::LocalSourceNode { }
+
+    private class ConnectionConstruction extends InstanceSource, DataFlow::CallCfgNode {
+      ConnectionConstruction() {
+        this = classRef().getACall()
+        or
+        this.(DataFlow::MethodCallNode).calls(Engine::instance(), ["begin", "connect"])
+        or
+        this.(DataFlow::MethodCallNode).calls(instance(), "connect")
+      }
+    }
+
+    /** Gets a reference to an instance of a SQLAlchemy Connection. */
+    private DataFlow::TypeTrackingNode instance(DataFlow::TypeTracker t) {
+      t.start() and
+      result instanceof InstanceSource
+      or
+      exists(DataFlow::TypeTracker t2 | result = instance(t2).track(t2, t))
+    }
+
+    /** Gets a reference to an instance of a SQLAlchemy Connection. */
+    DataFlow::Node instance() { instance(DataFlow::TypeTracker::end()).flowsTo(result) }
   }
 
   /**
-   * A call to `scalar` meant to execute an SQL expression
-   * See https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.Session.scalar and
-   * https://docs.sqlalchemy.org/en/14/core/connections.html?highlight=scalar#sqlalchemy.engine.Engine.scalar
+   * Provides models for the `sqlalchemy.orm.Session` class
+   *
+   * See
+   * - https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.Session
+   * - https://docs.sqlalchemy.org/en/14/orm/session_basics.html
    */
-  private class SqlAlchemyScalarCall extends DataFlow::CallCfgNode, SqlExecution::Range {
-    SqlAlchemyScalarCall() {
-      this =
-        [getSqlAlchemySessionInstance(), getSqlAlchemyEngineInstance()]
-            .getMember("scalar")
-            .getACall()
+  module Session {
+    /** Gets a reference to the `sqlalchemy.orm.Session` class. */
+    private API::Node classRef() {
+      result = API::moduleImport("sqlalchemy").getMember("orm").getMember("Session")
     }
 
-    override DataFlow::Node getSql() { result = this.getArg(0) }
+    /**
+     * A source of instances of `sqlalchemy.orm.Session`, extend this class to model new instances.
+     *
+     * This can include instantiations of the class, return values from function
+     * calls, or a special parameter that will be set when functions are called by an external
+     * library.
+     *
+     * Use the predicate `Session::instance()` to get references to instances of `sqlalchemy.orm.Session`.
+     */
+    abstract class InstanceSource extends DataFlow::LocalSourceNode { }
+
+    private class SessionConstruction extends InstanceSource, DataFlow::CallCfgNode {
+      SessionConstruction() {
+        this = classRef().getACall()
+        or
+        this =
+          API::moduleImport("sqlalchemy")
+              .getMember("orm")
+              .getMember("sessionmaker")
+              .getReturn()
+              .getACall()
+      }
+    }
+
+    /** Gets a reference to an instance of `sqlalchemy.orm.Session`. */
+    private DataFlow::TypeTrackingNode instance(DataFlow::TypeTracker t) {
+      t.start() and
+      result instanceof InstanceSource
+      or
+      exists(DataFlow::TypeTracker t2 | result = instance(t2).track(t2, t))
+    }
+
+    /** Gets a reference to an instance of `sqlalchemy.orm.Session`. */
+    DataFlow::Node instance() { instance(DataFlow::TypeTracker::end()).flowsTo(result) }
   }
 
   /**
-   * A call on a Query object
-   * See https://docs.sqlalchemy.org/en/14/orm/query.html?highlight=query#sqlalchemy.orm.Query
+   * A call to `execute` on a SQLAlchemy Engine, Connection, or Session.
+   * See
+   *  - https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.Engine.execute
+   *  - https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.Connection.execute
+   *  - https://docs.sqlalchemy.org/en/14/core/future.html#sqlalchemy.future.Connection.execute
+   *  - https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.Session.execute
    */
-  private class SqlAlchemyQueryCall extends DataFlow::CallCfgNode, SqlExecution::Range {
-    SqlAlchemyQueryCall() {
-      this =
-        getSqlAlchemyQueryInstance()
-            .getMember(any(SqlAlchemyVulnerableMethodNames methodName))
-            .getACall()
+  private class SqlAlchemyExecuteCall extends DataFlow::MethodCallNode, SqlExecution::Range {
+    SqlAlchemyExecuteCall() {
+      this.calls(Engine::instance(), "execute")
+      or
+      this.calls(Connection::instance(), "execute")
+      or
+      this.calls(Session::instance(), "execute")
     }
 
-    override DataFlow::Node getSql() { result = this.getArg(0) }
+    override DataFlow::Node getSql() { result in [this.getArg(0), this.getArgByName("statement")] }
   }
 
   /**
-   * This class represents a list of methods vulnerable to sql injection.
-   *
-   * See https://github.com/jty-team/codeql/pull/2#issue-611592361
+   * A call to `scalar` on a SQLAlchemy Engine, Connection, or Session.
+   * See
+   *  - https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.Engine.scalar
+   *  - https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.Connection.scalar
+   *  - https://docs.sqlalchemy.org/en/14/core/future.html#sqlalchemy.future.Connection.scalar
+   *  - https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.Session.scalar
    */
-  private class SqlAlchemyVulnerableMethodNames extends string {
-    SqlAlchemyVulnerableMethodNames() { this in ["filter", "filter_by", "group_by", "order_by"] }
+  private class SqlAlchemyScalarCall extends DataFlow::MethodCallNode, SqlExecution::Range {
+    SqlAlchemyScalarCall() {
+      this.calls(Engine::instance(), "scalar")
+      or
+      this.calls(Connection::instance(), "scalar")
+      or
+      this.calls(Session::instance(), "scalar")
+    }
+
+    override DataFlow::Node getSql() {
+      result in [this.getArg(0), this.getArgByName("statement"), this.getArgByName("object_")]
+    }
   }
 
   /**
@@ -146,3 +264,47 @@ private module SqlAlchemy {
     override DataFlow::Node getAnInput() { result = this.getArg(0) }
   }
 }
+
+private module OldModeling {
+  /**
+   * Returns an instantization of a SqlAlchemy Session object.
+   * See https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.Session and
+   * https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.sessionmaker
+   */
+  private API::Node getSqlAlchemySessionInstance() {
+    result = API::moduleImport("sqlalchemy.orm").getMember("Session").getReturn() or
+    result = API::moduleImport("sqlalchemy.orm").getMember("sessionmaker").getReturn().getReturn()
+  }
+
+  /**
+   * Returns an instantization of a SqlAlchemy Query object.
+   * See https://docs.sqlalchemy.org/en/14/orm/query.html?highlight=query#sqlalchemy.orm.Query
+   */
+  private API::Node getSqlAlchemyQueryInstance() {
+    result = getSqlAlchemySessionInstance().getMember("query").getReturn()
+  }
+
+  /**
+   * A call on a Query object
+   * See https://docs.sqlalchemy.org/en/14/orm/query.html?highlight=query#sqlalchemy.orm.Query
+   */
+  private class SqlAlchemyQueryCall extends DataFlow::CallCfgNode, SqlExecution::Range {
+    SqlAlchemyQueryCall() {
+      this =
+        getSqlAlchemyQueryInstance()
+            .getMember(any(SqlAlchemyVulnerableMethodNames methodName))
+            .getACall()
+    }
+
+    override DataFlow::Node getSql() { result = this.getArg(0) }
+  }
+
+  /**
+   * This class represents a list of methods vulnerable to sql injection.
+   *
+   * See https://github.com/jty-team/codeql/pull/2#issue-611592361
+   */
+  private class SqlAlchemyVulnerableMethodNames extends string {
+    SqlAlchemyVulnerableMethodNames() { this in ["filter", "filter_by", "group_by", "order_by"] }
+  }
+}