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;