Python: Add aiohttp.ClientSession model

tausbn · RasmusWL · commit 947aa099e031 · 2023-12-08T11:27:51.000+01:00
diff --git a/python/ql/lib/semmle/python/frameworks/Aiohttp.qll b/python/ql/lib/semmle/python/frameworks/Aiohttp.qll
@@ -14,6 +14,7 @@ private import semmle.python.frameworks.internal.SelfRefMixin
 private import semmle.python.frameworks.Multidict
 private import semmle.python.frameworks.Yarl
 private import semmle.python.frameworks.internal.InstanceTaintStepsHelper
+private import semmle.python.frameworks.data.ModelsAsData
 
 /**
  * INTERNAL: Do not use.
@@ -706,19 +707,23 @@ module AiohttpWebModel {
 }
 
 /**
+ * INTERNAL: Do not use.
+ *
  * Provides models for the web server part (`aiohttp.client`) of the `aiohttp` PyPI package.
  * See https://docs.aiohttp.org/en/stable/client.html
  */
-private module AiohttpClientModel {
+module AiohttpClientModel {
   /**
    * Provides models for the `aiohttp.ClientSession` class
    *
    * See https://docs.aiohttp.org/en/stable/client_reference.html#aiohttp.ClientSession.
    */
   module ClientSession {
     /** Gets a reference to the `aiohttp.ClientSession` class. */
-    private API::Node classRef() {
+    API::Node classRef() {
       result = API::moduleImport("aiohttp").getMember("ClientSession")
+      or
+      result = ModelOutput::getATypeNode("aiohttp.ClientSession~Subclass").getASubclass*()
     }
 
     /** Gets a reference to an instance of `aiohttp.ClientSession`. */
diff --git a/python/ql/src/meta/ClassHierarchy/Find.ql b/python/ql/src/meta/ClassHierarchy/Find.ql
@@ -16,6 +16,7 @@ private import semmle.python.frameworks.Stdlib
 private import semmle.python.frameworks.Requests
 private import semmle.python.frameworks.Starlette
 private import semmle.python.frameworks.ClickhouseDriver
+private import semmle.python.frameworks.Aiohttp
 import semmle.python.frameworks.data.internal.ApiGraphModelsExtensions as Extensions
 
 class FlaskViewClasses extends FindSubclassesSpec {
@@ -241,6 +242,14 @@ class ClickhouseClient extends FindSubclassesSpec {
   override API::Node getAlreadyModeledClass() { result = ClickhouseDriver::Client::subclassRef() }
 }
 
+class AiohttpSession extends FindSubclassesSpec {
+  AiohttpSession() { this = "aiohttp.ClientSession~Subclass" }
+
+  override API::Node getAlreadyModeledClass() {
+    result = AiohttpClientModel::ClientSession::classRef()
+  }
+}
+
 bindingset[fullyQualified]
 predicate fullyQualifiedToYamlFormat(string fullyQualified, string type2, string path) {
   exists(int firstDot | firstDot = fullyQualified.indexOf(".", 0, 0) |
