ATM: Extract training data

javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointCharacteristics.qll

@@ -45,6 +45,9 @@ abstract class EndpointCharacteristic extends string {
     EndpointType endpointClass, boolean isPositiveIndicator, float confidence
   );
 
+  /** Indicators with confidence at or above this threshold are considered to be high-confidence indicators. */
+  final float getHighConfidenceThreshold() { result = 0.8 }
+
   // The following are some confidence values that are used in practice by the subclasses. They are defined as named
   // constants here to make it easier to change them in the future.
   final float maximalConfidence() { result = 1.0 }

javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/DebugResultInclusion.ql

@@ -11,7 +11,7 @@
 
 import javascript
 import experimental.adaptivethreatmodeling.ATMConfig
-import extraction.ExtractEndpointData
+import extraction.ExtractEndpointDataTraining
 
 string getAReasonSinkExcluded(DataFlow::Node sinkCandidate, Query query) {
   query instanceof NosqlInjectionQuery and
@@ -33,7 +33,7 @@ string getDescriptionForAlertCandidate(
 ) {
   result = "excluded[reason=" + getAReasonSinkExcluded(sinkCandidate, query) + "]"
   or
-  getAtmCfg(query).isKnownSink(sinkCandidate) and
+  getDataFlowCfg(query).(AtmConfig).isKnownSink(sinkCandidate) and
   result = "excluded[reason=known-sink]"
   or
   not exists(getAReasonSinkExcluded(sinkCandidate, query)) and

javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/extraction/ExtractEndpointDataTraining.ql

@@ -4,23 +4,8 @@
  * Extracts training data we can use to train ML models for ML-powered queries.
  */
 
-import javascript
-import ExtractEndpointData as ExtractEndpointData
+private import ExtractEndpointDataTraining as ExtractEndpointDataTraining
 
-query predicate endpoints(
-  DataFlow::Node endpoint, string queryName, string key, string value, string valueType
-) {
-  ExtractEndpointData::endpoints(endpoint, queryName, key, value, valueType) and
-  // only select endpoints that are either Sink or NotASink
-  ExtractEndpointData::endpoints(endpoint, queryName, "sinkLabel", ["Sink", "NotASink"], "string") and
-  // do not select endpoints filtered out by end-to-end evaluation
-  ExtractEndpointData::endpoints(endpoint, queryName, "isExcludedFromEndToEndEvaluation", "false",
-    "boolean") and
-  // only select endpoints that can be part of a tainted flow
-  ExtractEndpointData::endpoints(endpoint, queryName, "isConstantExpression", "false", "boolean")
-}
+query predicate endpoints = ExtractEndpointDataTraining::reformattedTrainingEndpoints/5;
 
-query predicate tokenFeatures(DataFlow::Node endpoint, string featureName, string featureValue) {
-  endpoints(endpoint, _, _, _, _) and
-  ExtractEndpointData::tokenFeatures(endpoint, featureName, featureValue)
-}
+query predicate tokenFeatures = ExtractEndpointDataTraining::tokenFeatures/3;