From 0bf1ca7550d652032f400ad1257404b101bc0974 Mon Sep 17 00:00:00 2001
From: Rasmus Wriedt Larsen <rasmuswl@github.com>
Date: Fri, 15 Sep 2023 14:58:42 +0200
Subject: [PATCH] Python: Automodel for `WSGIServer`

---
 .../lib/semmle/python/frameworks/Stdlib.qll   | 34 +++++++++----------
 python/ql/src/meta/ClassHierarchy/Find.ql     |  9 +++++
 2 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/python/ql/lib/semmle/python/frameworks/Stdlib.qll b/python/ql/lib/semmle/python/frameworks/Stdlib.qll
index ee33f3f1210..da1b7758489 100644
--- a/python/ql/lib/semmle/python/frameworks/Stdlib.qll
+++ b/python/ql/lib/semmle/python/frameworks/Stdlib.qll
@@ -17,6 +17,7 @@ private import semmle.python.frameworks.internal.InstanceTaintStepsHelper
 // modeling split over multiple files to keep this file from becoming too big
 private import semmle.python.frameworks.Stdlib.Urllib
 private import semmle.python.frameworks.Stdlib.Urllib2
+private import semmle.python.frameworks.data.ModelsAsData
 
 /** Provides models for the Python standard library. */
 module Stdlib {
@@ -298,7 +299,7 @@ module Stdlib {
  * policy, and the code is not in a polished enough state that we want to do so -- at
  * least not without having convincing use-cases for it :)
  */
-private module StdlibPrivate {
+module StdlibPrivate {
   // ---------------------------------------------------------------------------
   // os
   // ---------------------------------------------------------------------------
@@ -2035,17 +2036,20 @@ private module StdlibPrivate {
   // wsgiref.simple_server
   // ---------------------------------------------------------------------------
   /** Provides models for the `wsgiref.simple_server` module. */
-  private module WsgirefSimpleServer {
+  module WsgirefSimpleServer {
+    API::Node subclassRef() {
+      result =
+        API::moduleImport("wsgiref")
+            .getMember("simple_server")
+            .getMember("WSGIServer")
+            .getASubclass*()
+      or
+      result =
+        ModelOutput::getATypeNode("wsgiref.simple_server.WSGIServer~Subclass").getASubclass*()
+    }
+
     class WsgiServerSubclass extends Class, SelfRefMixin {
-      WsgiServerSubclass() {
-        this.getParent() =
-          API::moduleImport("wsgiref")
-              .getMember("simple_server")
-              .getMember("WSGIServer")
-              .getASubclass*()
-              .asSource()
-              .asExpr()
-      }
+      WsgiServerSubclass() { this.getParent() = subclassRef().asSource().asExpr() }
     }
 
     /**
@@ -2062,13 +2066,7 @@ private module StdlibPrivate {
         exists(DataFlow::Node appArg, DataFlow::CallCfgNode setAppCall |
           (
             setAppCall =
-              API::moduleImport("wsgiref")
-                  .getMember("simple_server")
-                  .getMember("WSGIServer")
-                  .getASubclass*()
-                  .getReturn()
-                  .getMember("set_app")
-                  .getACall()
+              WsgirefSimpleServer::subclassRef().getReturn().getMember("set_app").getACall()
             or
             setAppCall
                 .(DataFlow::MethodCallNode)
diff --git a/python/ql/src/meta/ClassHierarchy/Find.ql b/python/ql/src/meta/ClassHierarchy/Find.ql
index 9fc62b68c1e..059a710bc6e 100644
--- a/python/ql/src/meta/ClassHierarchy/Find.ql
+++ b/python/ql/src/meta/ClassHierarchy/Find.ql
@@ -12,6 +12,7 @@ private import semmle.python.frameworks.Flask
 private import semmle.python.frameworks.FastApi
 private import semmle.python.frameworks.Django
 private import semmle.python.frameworks.Tornado
+private import semmle.python.frameworks.Stdlib
 import semmle.python.frameworks.data.internal.ApiGraphModelsExtensions as Extensions
 
 class FlaskViewClasses extends FindSubclassesSpec {
@@ -66,6 +67,14 @@ class TornadoRequestHandler extends FindSubclassesSpec {
   }
 }
 
+class WSGIServer extends FindSubclassesSpec {
+  WSGIServer() { this = "wsgiref.simple_server.WSGIServer~Subclass" }
+
+  override API::Node getAlreadyModeledClass() {
+    result = StdlibPrivate::WsgirefSimpleServer::subclassRef()
+  }
+}
+
 bindingset[fullyQualified]
 predicate fullyQualifiedToYamlFormat(string fullyQualified, string type2, string path) {
   exists(int firstDot | firstDot = fullyQualified.indexOf(".", 0, 0) |