Finish Sinks

Author: jorgectf

python/ql/src/experimental/Security/CWE-113/HeaderInjection.ql

@@ -31,7 +31,7 @@ class WerkzeugHeaderCall extends DataFlow::CallCfgNode {
   DataFlow::Node getHeaderInputNode() { result = this.getArg(1) }
 }
 
-class FlaskHeaderCall extends DataFlow::CallCfgNode {
+class FlaskHeaderCall extends DataFlow::Node {
   DataFlow::Node headerInputNode;
 
   FlaskHeaderCall() {
@@ -44,39 +44,108 @@ class FlaskHeaderCall extends DataFlow::CallCfgNode {
       responseMethod.getObject().getALocalSource() = headerInstance and
       sinkDeclaration.getATarget() = responseMethod.asExpr().getParentNode() and
       headerInputNode.asExpr() = sinkDeclaration.getValue() and
-      this.getFunction() = responseMethod
+      this.asExpr() = sinkDeclaration.getATarget()
     )
   }
 
   DataFlow::Node getHeaderInputNode() { result = headerInputNode }
 }
 
-class FlaskMakeResponse extends DataFlow::Node {
-  FlaskMakeResponse() {
+class FlaskMakeResponseCall extends DataFlow::Node {
+  DataFlow::Node headerInputNode;
+
+  FlaskMakeResponseCall() {
     exists(
       DataFlow::CallCfgNode headerInstance, DataFlow::AttrRead responseMethod,
       AssignStmt sinkDeclaration
     |
       headerInstance = API::moduleImport("flask").getMember("make_response").getACall() and
       responseMethod.getAttributeName() = "headers" and
       responseMethod.getObject().getALocalSource() = headerInstance and
-      (
-        sinkDeclaration.getATarget() = responseMethod.asExpr().getParentNode() and
-        this.asExpr() = sinkDeclaration.getValue()
-      )
-      //or
-      //extendMethod.getAttributeName() = "extend" and
-      //extendMethod.getObject().getALocalSource() = responseMethod and
-      //this = extendMethod.(DataFlow::CallCfgNode).getArg(0)
+      sinkDeclaration.getATarget() = responseMethod.asExpr().getParentNode() and
+      this.asExpr() = sinkDeclaration.getATarget() and
+      headerInputNode.asExpr() = sinkDeclaration.getValue()
+    )
+  }
+
+  DataFlow::Node getHeaderInputNode() { result = headerInputNode }
+}
+
+class FlaskMakeResponseExtendCall extends DataFlow::CallCfgNode {
+  DataFlow::Node headerInputNode;
+
+  FlaskMakeResponseExtendCall() {
+    exists(
+      DataFlow::CallCfgNode headerInstance, DataFlow::AttrRead responseMethod,
+      DataFlow::AttrRead extendMethod
+    |
+      headerInstance = API::moduleImport("flask").getMember("make_response").getACall() and
+      responseMethod.getAttributeName() = "headers" and
+      responseMethod.getObject().getALocalSource() = headerInstance and
+      extendMethod.getAttributeName() = "extend" and
+      extendMethod.getObject().getALocalSource() = responseMethod and
+      this.getFunction() = extendMethod and
+      headerInputNode = this.getArg(0)
     )
   }
+
+  DataFlow::Node getHeaderInputNode() { result = headerInputNode }
+}
+
+class FlaskResponseArg extends DataFlow::CallCfgNode {
+  DataFlow::Node headerInputNode;
+
+  FlaskResponseArg() {
+    this = API::moduleImport("flask").getMember("Response").getACall() and
+    headerInputNode = this.getArgByName("headers")
+  }
+
+  DataFlow::Node getHeaderInputNode() { result = headerInputNode }
+}
+
+class DjangoResponseSetItemCall extends DataFlow::CallCfgNode {
+  DjangoResponseSetItemCall() {
+    exists(DataFlow::AttrRead setItemMethod |
+      this.getFunction() = setItemMethod and
+      setItemMethod.getObject().getALocalSource() =
+        API::moduleImport("django").getMember("http").getMember("HttpResponse").getACall() and
+      setItemMethod.getAttributeName() = "__setitem__"
+    )
+  }
+
+  DataFlow::Node getHeaderInputNode() { result = this.getArg(1) }
+}
+
+class DjangoResponseAssignCall extends DataFlow::Node {
+  DataFlow::Node headerInputNode;
+
+  DjangoResponseAssignCall() {
+    exists(
+      DataFlow::CallCfgNode headerInstance, Subscript responseMethod, DataFlow::Node responseToNode,
+      AssignStmt sinkDeclaration
+    |
+      headerInstance =
+        API::moduleImport("django").getMember("http").getMember("HttpResponse").getACall() and
+      responseMethod.getValue() = responseToNode.asExpr() and
+      responseToNode.getALocalSource().asExpr() = headerInstance.asExpr() and
+      sinkDeclaration.getATarget() = responseMethod and
+      this.asExpr() = sinkDeclaration.getATarget() and
+      headerInputNode.asExpr() = sinkDeclaration.getValue()
+    )
+  }
+
+  DataFlow::Node getHeaderInputNode() { result = headerInputNode }
 }
 
 class HeaderInjectionSink extends DataFlow::Node {
   HeaderInjectionSink() {
-    this instanceof WerkzeugHeaderCall or
-    this instanceof FlaskHeaderCall or
-    this instanceof FlaskMakeResponse
+    this = any(WerkzeugHeaderCall a).getHeaderInputNode() or
+    this = any(FlaskHeaderCall a).getHeaderInputNode() or
+    this = any(FlaskMakeResponseCall a).getHeaderInputNode() or
+    this = any(FlaskMakeResponseExtendCall a).getHeaderInputNode() or
+    this = any(FlaskResponseArg a).getHeaderInputNode() or
+    this = any(DjangoResponseSetItemCall a).getHeaderInputNode() or
+    this = any(DjangoResponseAssignCall a).getHeaderInputNode()
   }
 }