Python: Add rest_framework model

tausbn · RasmusWL · commit 422733f32a0e · 2023-12-08T11:27:51.000+01:00
diff --git a/python/ql/lib/semmle/python/frameworks/RestFramework.qll b/python/ql/lib/semmle/python/frameworks/RestFramework.qll
@@ -16,6 +16,7 @@ private import semmle.python.ApiGraphs
 private import semmle.python.frameworks.internal.InstanceTaintStepsHelper
 private import semmle.python.frameworks.Django
 private import semmle.python.frameworks.Stdlib
+private import semmle.python.frameworks.data.ModelsAsData
 
 /**
  * INTERNAL: Do not use.
@@ -27,7 +28,7 @@ private import semmle.python.frameworks.Stdlib
  * - https://www.django-rest-framework.org/
  * - https://pypi.org/project/djangorestframework/
  */
-private module RestFramework {
+module RestFramework {
   // ---------------------------------------------------------------------------
   // rest_framework.views.APIView handling
   // ---------------------------------------------------------------------------
@@ -215,8 +216,10 @@ private module RestFramework {
    */
   module Request {
     /** Gets a reference to the `rest_framework.request.Request` class. */
-    private API::Node classRef() {
+    API::Node classRef() {
       result = API::moduleImport("rest_framework").getMember("request").getMember("Request")
+      or
+      result = ModelOutput::getATypeNode("rest_framework.request.Request~Subclass").getASubclass*()
     }
 
     /**
@@ -299,8 +302,11 @@ private module RestFramework {
    */
   module Response {
     /** Gets a reference to the `rest_framework.response.Response` class. */
-    private API::Node classRef() {
+    API::Node classRef() {
       result = API::moduleImport("rest_framework").getMember("response").getMember("Response")
+      or
+      result =
+        ModelOutput::getATypeNode("rest_framework.response.Response~Subclass").getASubclass*()
     }
 
     /** A direct instantiation of `rest_framework.response.Response`. */
diff --git a/python/ql/src/meta/ClassHierarchy/Find.ql b/python/ql/src/meta/ClassHierarchy/Find.ql
@@ -23,6 +23,7 @@ private import semmle.python.frameworks.Invoke
 private import semmle.python.frameworks.MarkupSafe
 private import semmle.python.frameworks.Multidict
 private import semmle.python.frameworks.Pycurl
+private import semmle.python.frameworks.RestFramework
 import semmle.python.frameworks.data.internal.ApiGraphModelsExtensions as Extensions
 
 class FlaskViewClasses extends FindSubclassesSpec {
@@ -324,6 +325,18 @@ class PyCurl extends FindSubclassesSpec {
   override API::Node getAlreadyModeledClass() { result = Pycurl::Curl::classRef() }
 }
 
+class RestFrameworkRequest extends FindSubclassesSpec {
+  RestFrameworkRequest() { this = "rest_framework.request.Request~Subclass" }
+
+  override API::Node getAlreadyModeledClass() { result = RestFramework::Request::classRef() }
+}
+
+class RestFrameworkResponse extends FindSubclassesSpec {
+  RestFrameworkResponse() { this = "rest_framework.response.Response~Subclass" }
+
+  override API::Node getAlreadyModeledClass() { result = RestFramework::Response::classRef() }
+}
+
 bindingset[fullyQualified]
 predicate fullyQualifiedToYamlFormat(string fullyQualified, string type2, string path) {
   exists(int firstDot | firstDot = fullyQualified.indexOf(".", 0, 0) |
