Python: model aiomysql

docs/codeql/support/reusables/frameworks.rst

@@ -183,6 +183,7 @@ Python built-in support
    pydantic, Utility library
    yarl, Utility library
    aioch, Database
+   aiomysql, Database
    asyncpg, Database
    clickhouse-driver, Database
    mysql-connector-python, Database

python/change-notes/2021-10-11-model-aiomysql.md

@@ -0,0 +1,2 @@
+lgtm,codescanning
+* Added modeling of `aiomysql` for sinks executing SQL

python/ql/lib/semmle/python/Frameworks.qll

@@ -6,6 +6,7 @@
 // `docs/codeql/support/reusables/frameworks.rst`
 private import semmle.python.frameworks.Aioch
 private import semmle.python.frameworks.Aiohttp
+private import semmle.python.frameworks.Aiomysql
 private import semmle.python.frameworks.Asyncpg
 private import semmle.python.frameworks.ClickhouseDriver
 private import semmle.python.frameworks.Cryptodome

python/ql/lib/semmle/python/frameworks/Aiomysql.qll

@@ -0,0 +1,145 @@
+/**
+ * Provides classes modeling security-relevant aspects of the `aiomysql` PyPI package.
+ * See
+ * - https://aiomysql.readthedocs.io/en/stable/index.html
+ * - https://pypi.org/project/aiomysql/
+ */
+
+private import python
+private import semmle.python.dataflow.new.DataFlow
+private import semmle.python.Concepts
+private import semmle.python.ApiGraphs
+
+/** Provides models for the `aiomysql` PyPI package. */
+private module Aiomysql {
+  private import semmle.python.internal.Awaited
+
+  /**
+   * A `ConectionPool` is created when the result of `aiomysql.create_pool()` is awaited.
+   * See https://aiomysql.readthedocs.io/en/stable/core.html#pool
+   */
+  API::Node connectionPool() {
+    result = API::moduleImport("aiomysql").getMember("create_pool").getReturn().getAwaited()
+  }
+
+  /**
+   * A `Connection` is created when
+   * - the result of `aiomysql.connect()` is awaited.
+   * - the result of calling `aquire` on a `ConnectionPool` is awaited.
+   * See https://aiomysql.readthedocs.io/en/stable/core.html#connection
+   */
+  API::Node connection() {
+    result = API::moduleImport("aiomysql").getMember("connect").getReturn().getAwaited()
+    or
+    result = connectionPool().getMember("acquire").getReturn().getAwaited()
+  }
+
+  /**
+   * A `Cursor` is created when
+   * - the result of calling `cursor` on a `ConnectionPool` is awaited.
+   * - the result of calling `cursor` on a `Connection` is awaited.
+   * See https://aiomysql.readthedocs.io/en/stable/core.html#cursor
+   */
+  API::Node cursor() {
+    result = connectionPool().getMember("cursor").getReturn().getAwaited()
+    or
+    result = connection().getMember("cursor").getReturn().getAwaited()
+  }
+
+  /**
+   * Calling `execute` on a `Cursor` constructs a query.
+   * See https://aiomysql.readthedocs.io/en/stable/core.html#aiomysql.Cursor.execute
+   */
+  class CursorExecuteCall extends SqlConstruction::Range, DataFlow::CallCfgNode {
+    CursorExecuteCall() { this = cursor().getMember("execute").getACall() }
+
+    override DataFlow::Node getSql() { result in [this.getArg(0), this.getArgByName("operation")] }
+  }
+
+  /**
+   * This is only needed to connect the argument to the execute call with the subsequnt awaiting.
+   * It should be obsolete once we have `API::CallNode` available.
+   */
+  private DataFlow::TypeTrackingNode cursorExecuteCall(DataFlow::TypeTracker t, DataFlow::Node sql) {
+    // cursor created from connection
+    t.start() and
+    sql = result.(CursorExecuteCall).getSql()
+    or
+    exists(DataFlow::TypeTracker t2 | result = cursorExecuteCall(t2, sql).track(t2, t))
+  }
+
+  DataFlow::Node cursorExecuteCall(DataFlow::Node sql) {
+    cursorExecuteCall(DataFlow::TypeTracker::end(), sql).flowsTo(result)
+  }
+
+  /**
+   * Awaiting the result of calling `execute` executes the query.
+   * See https://aiomysql.readthedocs.io/en/stable/core.html#aiomysql.Cursor.execute
+   */
+  class AwaitedCursorExecuteCall extends SqlExecution::Range {
+    DataFlow::Node sql;
+
+    AwaitedCursorExecuteCall() { this = awaited(cursorExecuteCall(sql)) }
+
+    override DataFlow::Node getSql() { result = sql }
+  }
+
+  /**
+   * An `Engine` is created when the result of calling `aiomysql.sa.create_engine` is awaited.
+   * See https://aiomysql.readthedocs.io/en/stable/sa.html#engine
+   */
+  API::Node engine() {
+    result =
+      API::moduleImport("aiomysql")
+          .getMember("sa")
+          .getMember("create_engine")
+          .getReturn()
+          .getAwaited()
+  }
+
+  /**
+   * A `SAConnection` is created when the result of calling `aquire` on an `Engine` is awaited.
+   * See https://aiomysql.readthedocs.io/en/stable/sa.html#connection
+   */
+  API::Node saConnection() { result = engine().getMember("acquire").getReturn().getAwaited() }
+
+  /**
+   * Calling `execute` on a `SAConnection` constructs a query.
+   * See https://aiomysql.readthedocs.io/en/stable/sa.html#aiomysql.sa.SAConnection.execute
+   */
+  class SAConnectionExecuteCall extends SqlConstruction::Range, DataFlow::CallCfgNode {
+    SAConnectionExecuteCall() { this = saConnection().getMember("execute").getACall() }
+
+    override DataFlow::Node getSql() { result in [this.getArg(0), this.getArgByName("query")] }
+  }
+
+  /**
+   * This is only needed to connect the argument to the execute call with the subsequnt awaiting.
+   * It should be obsolete once we have `API::CallNode` available.
+   */
+  private DataFlow::TypeTrackingNode saConnectionExecuteCall(
+    DataFlow::TypeTracker t, DataFlow::Node sql
+  ) {
+    // saConnection created from engine
+    t.start() and
+    sql = result.(SAConnectionExecuteCall).getSql()
+    or
+    exists(DataFlow::TypeTracker t2 | result = saConnectionExecuteCall(t2, sql).track(t2, t))
+  }
+
+  DataFlow::Node saConnectionExecuteCall(DataFlow::Node sql) {
+    saConnectionExecuteCall(DataFlow::TypeTracker::end(), sql).flowsTo(result)
+  }
+
+  /**
+   * Awaiting the result of calling `execute` executes the query.
+   * See https://aiomysql.readthedocs.io/en/stable/sa.html#aiomysql.sa.SAConnection.execute
+   */
+  class AwaitedSAConnectionExecuteCall extends SqlExecution::Range {
+    DataFlow::Node sql;
+
+    AwaitedSAConnectionExecuteCall() { this = awaited(saConnectionExecuteCall(sql)) }
+
+    override DataFlow::Node getSql() { result = sql }
+  }
+}

python/ql/test/library-tests/frameworks/aiomysql/test.py

@@ -5,24 +5,24 @@ async def test_cursor():
     # Create connection directly
     conn = await aiomysql.connect()
     cur = await conn.cursor()
-    await cur.execute("sql")  # $ MISSING: getSql="sql" constructedSql="sql"
+    await cur.execute("sql")  # $ getSql="sql" constructedSql="sql"
 
     # Create connection via pool
     async with aiomysql.create_pool() as pool:
         # Create Cursor via Connection
         async with pool.acquire() as conn:
             async with conn.cursor() as cur:
-                await cur.execute("sql")  # $ MISSING: getSql="sql" constructedSql="sql"
+                await cur.execute("sql")  # $ getSql="sql" constructedSql="sql"
 
         # Create Cursor directly
         async with pool.cursor() as cur:
-            await cur.execute("sql")  # $ MISSING: getSql="sql" constructedSql="sql"
+            await cur.execute("sql")  # $ getSql="sql" constructedSql="sql"
 
     # variants using as few `async with` as possible
     pool = await aiomysql.create_pool()
     conn = await pool.acquire()
     cur = await conn.cursor()
-    await cur.execute("sql")  # $ MISSING: getSql="sql" constructedSql="sql"
+    await cur.execute("sql")  # $ getSql="sql" constructedSql="sql"
 
 # Test SQLAlchemy integration
 from aiomysql.sa import create_engine
@@ -30,4 +30,4 @@ from aiomysql.sa import create_engine
 async def test_engine():
     engine = await create_engine()
     conn = await engine.acquire()
-    await conn.execute("sql")  # $ MISSING: getSql="sql" constructedSql="sql"
+    await conn.execute("sql")  # $ getSql="sql" constructedSql="sql"