Python: Add support for threat models

python/ql/lib/qlpack.yml

@@ -9,6 +9,7 @@ dependencies:
   codeql/dataflow: ${workspace}
   codeql/mad: ${workspace}
   codeql/regex: ${workspace}
+  codeql/threat-models: ${workspace}
   codeql/tutorial: ${workspace}
   codeql/util: ${workspace}
   codeql/xml: ${workspace}

python/ql/lib/semmle/python/Concepts.qll

@@ -10,6 +10,54 @@ private import semmle.python.dataflow.new.RemoteFlowSources
 private import semmle.python.dataflow.new.TaintTracking
 private import semmle.python.Frameworks
 private import semmle.python.security.internal.EncryptionKeySizes
+private import codeql.threatmodels.ThreatModels
+
+/**
+ * A data flow source, for a specific threat-model.
+ *
+ * Extend this class to refine existing API models. If you want to model new APIs,
+ * extend `ThreatModelSource::Range` instead.
+ */
+class ThreatModelSource extends DataFlow::Node instanceof ThreatModelSource::Range {
+  /**
+   * Gets a string that represents the source kind with respect to threat modeling.
+   */
+  string getThreatModel() { result = super.getThreatModel() }
+
+  /** Gets a string that describes the type of this threat-model source. */
+  string getSourceType() { result = super.getSourceType() }
+}
+
+/** Provides a class for modeling new sources for specific threat-models. */
+module ThreatModelSource {
+  /**
+   * A data flow source, for a specific threat-model.
+   *
+   * Extend this class to model new APIs. If you want to refine existing API models,
+   * extend `ThreatModelSource` instead.
+   */
+  abstract class Range extends DataFlow::Node {
+    /**
+     * Gets a string that represents the source kind with respect to threat modeling.
+     */
+    abstract string getThreatModel();
+
+    /** Gets a string that describes the type of this threat-model source. */
+    abstract string getSourceType();
+  }
+}
+
+/**
+ * A data flow source that is enabled in the current threat model configuration.
+ */
+class ActiveThreatModelSource extends DataFlow::Node {
+  ActiveThreatModelSource() {
+    exists(string kind |
+      currentThreatModel(kind) and
+      this.(ThreatModelSource).getThreatModel() = kind
+    )
+  }
+}
 
 /**
  * A data-flow node that executes an operating system command,

python/ql/lib/semmle/python/dataflow/new/RemoteFlowSources.qll

@@ -15,10 +15,7 @@ private import semmle.python.Concepts
  * Extend this class to refine existing API models. If you want to model new APIs,
  * extend `RemoteFlowSource::Range` instead.
  */
-class RemoteFlowSource extends DataFlow::Node instanceof RemoteFlowSource::Range {
-  /** Gets a string that describes the type of this remote flow source. */
-  string getSourceType() { result = super.getSourceType() }
-}
+class RemoteFlowSource extends ThreatModelSource instanceof RemoteFlowSource::Range { }
 
 /** Provides a class for modeling new sources of remote user input. */
 module RemoteFlowSource {
@@ -28,8 +25,7 @@ module RemoteFlowSource {
    * Extend this class to model new APIs. If you want to refine existing API models,
    * extend `RemoteFlowSource` instead.
    */
-  abstract class Range extends DataFlow::Node {
-    /** Gets a string that describes the type of this remote flow source. */
-    abstract string getSourceType();
+  abstract class Range extends ThreatModelSource::Range {
+    override string getThreatModel() { result = "remote" }
   }
 }

python/ql/lib/semmle/python/frameworks/data/ModelsAsData.qll

@@ -18,14 +18,19 @@ private import semmle.python.dataflow.new.RemoteFlowSources
 private import semmle.python.dataflow.new.DataFlow
 private import semmle.python.ApiGraphs
 private import semmle.python.dataflow.new.FlowSummary
+private import semmle.python.Concepts
 
 /**
- * A remote flow source originating from a CSV source row.
+ * A threat-model flow source originating from a data extension.
  */
-private class RemoteFlowSourceFromCsv extends RemoteFlowSource::Range {
-  RemoteFlowSourceFromCsv() { this = ModelOutput::getASourceNode("remote").asSource() }
+private class ThreatModelSourceFromDataExtension extends ThreatModelSource::Range {
+  ThreatModelSourceFromDataExtension() { this = ModelOutput::getASourceNode(_).asSource() }
 
-  override string getSourceType() { result = "Remote flow (from model)" }
+  override string getThreatModel() { this = ModelOutput::getASourceNode(result).asSource() }
+
+  override string getSourceType() {
+    result = "Source node (" + this.getThreatModel() + ") [from data-extension]"
+  }
 }
 
 private class SummarizedCallableFromModel extends SummarizedCallable {