Skip to content

Commit 57e7bfb

Browse files
committed
Python: model aiomysql
1 parent 047cff0 commit 57e7bfb

File tree

5 files changed

+154
-5
lines changed

5 files changed

+154
-5
lines changed

docs/codeql/support/reusables/frameworks.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,7 @@ Python built-in support
183183
pydantic, Utility library
184184
yarl, Utility library
185185
aioch, Database
186+
aiomysql, Database
186187
asyncpg, Database
187188
clickhouse-driver, Database
188189
mysql-connector-python, Database
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
lgtm,codescanning
2+
* Added modeling of `aiomysql` for sinks executing SQL

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
// `docs/codeql/support/reusables/frameworks.rst`
77
private import semmle.python.frameworks.Aioch
88
private import semmle.python.frameworks.Aiohttp
9+
private import semmle.python.frameworks.Aiomysql
910
private import semmle.python.frameworks.Asyncpg
1011
private import semmle.python.frameworks.ClickhouseDriver
1112
private import semmle.python.frameworks.Cryptodome
Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
/**
2+
* Provides classes modeling security-relevant aspects of the `aiomysql` PyPI package.
3+
* See
4+
* - https://aiomysql.readthedocs.io/en/stable/index.html
5+
* - https://pypi.org/project/aiomysql/
6+
*/
7+
8+
private import python
9+
private import semmle.python.dataflow.new.DataFlow
10+
private import semmle.python.Concepts
11+
private import semmle.python.ApiGraphs
12+
13+
/** Provides models for the `aiomysql` PyPI package. */
14+
private module Aiomysql {
15+
private import semmle.python.internal.Awaited
16+
17+
/**
18+
* A `ConectionPool` is created when the result of `aiomysql.create_pool()` is awaited.
19+
* See https://aiomysql.readthedocs.io/en/stable/core.html#pool
20+
*/
21+
API::Node connectionPool() {
22+
result = API::moduleImport("aiomysql").getMember("create_pool").getReturn().getAwaited()
23+
}
24+
25+
/**
26+
* A `Connection` is created when
27+
* - the result of `aiomysql.connect()` is awaited.
28+
* - the result of calling `aquire` on a `ConnectionPool` is awaited.
29+
* See https://aiomysql.readthedocs.io/en/stable/core.html#connection
30+
*/
31+
API::Node connection() {
32+
result = API::moduleImport("aiomysql").getMember("connect").getReturn().getAwaited()
33+
or
34+
result = connectionPool().getMember("acquire").getReturn().getAwaited()
35+
}
36+
37+
/**
38+
* A `Cursor` is created when
39+
* - the result of calling `cursor` on a `ConnectionPool` is awaited.
40+
* - the result of calling `cursor` on a `Connection` is awaited.
41+
* See https://aiomysql.readthedocs.io/en/stable/core.html#cursor
42+
*/
43+
API::Node cursor() {
44+
result = connectionPool().getMember("cursor").getReturn().getAwaited()
45+
or
46+
result = connection().getMember("cursor").getReturn().getAwaited()
47+
}
48+
49+
/**
50+
* Calling `execute` on a `Cursor` constructs a query.
51+
* See https://aiomysql.readthedocs.io/en/stable/core.html#aiomysql.Cursor.execute
52+
*/
53+
class CursorExecuteCall extends SqlConstruction::Range, DataFlow::CallCfgNode {
54+
CursorExecuteCall() { this = cursor().getMember("execute").getACall() }
55+
56+
override DataFlow::Node getSql() { result in [this.getArg(0), this.getArgByName("operation")] }
57+
}
58+
59+
/**
60+
* This is only needed to connect the argument to the execute call with the subsequnt awaiting.
61+
* It should be obsolete once we have `API::CallNode` available.
62+
*/
63+
private DataFlow::TypeTrackingNode cursorExecuteCall(DataFlow::TypeTracker t, DataFlow::Node sql) {
64+
// cursor created from connection
65+
t.start() and
66+
sql = result.(CursorExecuteCall).getSql()
67+
or
68+
exists(DataFlow::TypeTracker t2 | result = cursorExecuteCall(t2, sql).track(t2, t))
69+
}
70+
71+
DataFlow::Node cursorExecuteCall(DataFlow::Node sql) {
72+
cursorExecuteCall(DataFlow::TypeTracker::end(), sql).flowsTo(result)
73+
}
74+
75+
/**
76+
* Awaiting the result of calling `execute` executes the query.
77+
* See https://aiomysql.readthedocs.io/en/stable/core.html#aiomysql.Cursor.execute
78+
*/
79+
class AwaitedCursorExecuteCall extends SqlExecution::Range {
80+
DataFlow::Node sql;
81+
82+
AwaitedCursorExecuteCall() { this = awaited(cursorExecuteCall(sql)) }
83+
84+
override DataFlow::Node getSql() { result = sql }
85+
}
86+
87+
/**
88+
* An `Engine` is created when the result of calling `aiomysql.sa.create_engine` is awaited.
89+
* See https://aiomysql.readthedocs.io/en/stable/sa.html#engine
90+
*/
91+
API::Node engine() {
92+
result =
93+
API::moduleImport("aiomysql")
94+
.getMember("sa")
95+
.getMember("create_engine")
96+
.getReturn()
97+
.getAwaited()
98+
}
99+
100+
/**
101+
* A `SAConnection` is created when the result of calling `aquire` on an `Engine` is awaited.
102+
* See https://aiomysql.readthedocs.io/en/stable/sa.html#connection
103+
*/
104+
API::Node saConnection() { result = engine().getMember("acquire").getReturn().getAwaited() }
105+
106+
/**
107+
* Calling `execute` on a `SAConnection` constructs a query.
108+
* See https://aiomysql.readthedocs.io/en/stable/sa.html#aiomysql.sa.SAConnection.execute
109+
*/
110+
class SAConnectionExecuteCall extends SqlConstruction::Range, DataFlow::CallCfgNode {
111+
SAConnectionExecuteCall() { this = saConnection().getMember("execute").getACall() }
112+
113+
override DataFlow::Node getSql() { result in [this.getArg(0), this.getArgByName("query")] }
114+
}
115+
116+
/**
117+
* This is only needed to connect the argument to the execute call with the subsequnt awaiting.
118+
* It should be obsolete once we have `API::CallNode` available.
119+
*/
120+
private DataFlow::TypeTrackingNode saConnectionExecuteCall(
121+
DataFlow::TypeTracker t, DataFlow::Node sql
122+
) {
123+
// saConnection created from engine
124+
t.start() and
125+
sql = result.(SAConnectionExecuteCall).getSql()
126+
or
127+
exists(DataFlow::TypeTracker t2 | result = saConnectionExecuteCall(t2, sql).track(t2, t))
128+
}
129+
130+
DataFlow::Node saConnectionExecuteCall(DataFlow::Node sql) {
131+
saConnectionExecuteCall(DataFlow::TypeTracker::end(), sql).flowsTo(result)
132+
}
133+
134+
/**
135+
* Awaiting the result of calling `execute` executes the query.
136+
* See https://aiomysql.readthedocs.io/en/stable/sa.html#aiomysql.sa.SAConnection.execute
137+
*/
138+
class AwaitedSAConnectionExecuteCall extends SqlExecution::Range {
139+
DataFlow::Node sql;
140+
141+
AwaitedSAConnectionExecuteCall() { this = awaited(saConnectionExecuteCall(sql)) }
142+
143+
override DataFlow::Node getSql() { result = sql }
144+
}
145+
}

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,29 +5,29 @@ async def test_cursor():
55
# Create connection directly
66
conn = await aiomysql.connect()
77
cur = await conn.cursor()
8-
await cur.execute("sql") # $ MISSING: getSql="sql" constructedSql="sql"
8+
await cur.execute("sql") # $ getSql="sql" constructedSql="sql"
99

1010
# Create connection via pool
1111
async with aiomysql.create_pool() as pool:
1212
# Create Cursor via Connection
1313
async with pool.acquire() as conn:
1414
async with conn.cursor() as cur:
15-
await cur.execute("sql") # $ MISSING: getSql="sql" constructedSql="sql"
15+
await cur.execute("sql") # $ getSql="sql" constructedSql="sql"
1616

1717
# Create Cursor directly
1818
async with pool.cursor() as cur:
19-
await cur.execute("sql") # $ MISSING: getSql="sql" constructedSql="sql"
19+
await cur.execute("sql") # $ getSql="sql" constructedSql="sql"
2020

2121
# variants using as few `async with` as possible
2222
pool = await aiomysql.create_pool()
2323
conn = await pool.acquire()
2424
cur = await conn.cursor()
25-
await cur.execute("sql") # $ MISSING: getSql="sql" constructedSql="sql"
25+
await cur.execute("sql") # $ getSql="sql" constructedSql="sql"
2626

2727
# Test SQLAlchemy integration
2828
from aiomysql.sa import create_engine
2929

3030
async def test_engine():
3131
engine = await create_engine()
3232
conn = await engine.acquire()
33-
await conn.execute("sql") # $ MISSING: getSql="sql" constructedSql="sql"
33+
await conn.execute("sql") # $ getSql="sql" constructedSql="sql"

0 commit comments

Comments
 (0)