Python: Add basic support for database threat-model

Author: RasmusWL

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

@@ -81,6 +81,24 @@ module PEP249 {
     }
   }
 
+  /** A call to a method that fetches rows from a previous execution. */
+  private class FetchMethodCall extends ThreatModelSource::Range, API::CallNode {
+    FetchMethodCall() {
+      exists(API::Node start |
+        start instanceof DatabaseCursor or start instanceof DatabaseConnection
+      |
+        // note: since we can't currently provide accesspaths for sources, these are all
+        // lumped together, although clearly the fetchmany/fetchall returns a
+        // list/iterable with rows.
+        this = start.getMember(["fetchone", "fetchmany", "fetchall"]).getACall()
+      )
+    }
+
+    override string getThreatModel() { result = "database" }
+
+    override string getSourceType() { result = "cursor.fetch*()" }
+  }
+
   // ---------------------------------------------------------------------------
   // asyncio implementations
   // ---------------------------------------------------------------------------

python/ql/test/library-tests/frameworks/psycopg/pep249.py

@@ -12,3 +12,24 @@
     with conn.cursor() as cursor:
         cursor.execute("some sql", (42,))  # $ getSql="some sql"
         cursor.executemany("some sql", [(42,)])  # $ getSql="some sql"
+
+
+        ### test of threat-model sources
+        row = cursor.fetchone() # $ threatModelSource[database]=cursor.fetchone()
+        rows_many = cursor.fetchmany(10) # $ threatModelSource[database]=cursor.fetchmany(..)
+        rows_all = cursor.fetchall() # $ threatModelSource[database]=cursor.fetchall()
+
+        ensure_tainted(
+            row[0],  # $ tainted
+            rows_many[0][0],  # $ tainted
+            rows_all[0][0],  # $ tainted
+
+            # pretending we created cursor to return dictionary results
+            row["column"],  # $ tainted
+            rows_many[0]["column"],  # $ tainted
+            rows_all[0]["column"],  # $ tainted
+        )
+        for row in rows_many:
+            ensure_tainted(row[0], row["column"]) # $ tainted
+        for row in rows_all:
+            ensure_tainted(row[0], row["column"]) # tainted