Merge pull request #2421 from github/koesie10/data-extension-editor-provenance

Add provenance support to data extensions editor

Author: koesie10

extensions/ql-vscode/src/data-extensions-editor/auto-model.ts

@@ -110,9 +110,10 @@ export function parsePredictedClassifications(
       // For now, model any method for which none of its arguments are modeled as sinks as neutral
       modeledMethods[signature] = {
         type: "neutral",
-        kind: "",
+        kind: "summary",
         input: "",
         output: "",
+        provenance: "ai-generated",
       };
       continue;
     }
@@ -129,6 +130,7 @@ export function parsePredictedClassifications(
       kind: sink.classification?.kind ?? "",
       input: sink.input ?? "",
       output: sink.output ?? "",
+      provenance: "ai-generated",
     };
   }
 

extensions/ql-vscode/src/data-extensions-editor/modeled-method.ts

@@ -5,11 +5,24 @@ export type ModeledMethodType =
   | "summary"
   | "neutral";
 
+export type Provenance =
+  // Generated by the dataflow model
+  | "df-generated"
+  // Generated by the dataflow model and manually edited
+  | "df-manual"
+  // Generated by the auto-model
+  | "ai-generated"
+  // Generated by the auto-model and manually edited
+  | "ai-manual"
+  // Entered by the user in the editor manually
+  | "manual";
+
 export type ModeledMethod = {
   type: ModeledMethodType;
   input: string;
   output: string;
   kind: string;
+  provenance: Provenance;
 };
 
 export type ModeledMethodWithSignature = {

extensions/ql-vscode/src/data-extensions-editor/predicates.ts

@@ -3,6 +3,7 @@ import {
   ModeledMethod,
   ModeledMethodType,
   ModeledMethodWithSignature,
+  Provenance,
 } from "./modeled-method";
 
 export type ExternalApiUsageByType = {
@@ -43,7 +44,7 @@ export const extensiblePredicateDefinitions: Record<
       "",
       method.modeledMethod.output,
       method.modeledMethod.kind,
-      "manual",
+      method.modeledMethod.provenance,
     ],
     readModeledMethod: (row) => ({
       signature: readRowToMethod(row),
@@ -52,6 +53,7 @@ export const extensiblePredicateDefinitions: Record<
         input: "",
         output: row[6] as string,
         kind: row[7] as string,
+        provenance: row[8] as Provenance,
       },
     }),
     supportedKinds: ["remote"],
@@ -71,7 +73,7 @@ export const extensiblePredicateDefinitions: Record<
       "",
       method.modeledMethod.input,
       method.modeledMethod.kind,
-      "manual",
+      method.modeledMethod.provenance,
     ],
     readModeledMethod: (row) => ({
       signature: readRowToMethod(row),
@@ -80,6 +82,7 @@ export const extensiblePredicateDefinitions: Record<
         input: row[6] as string,
         output: "",
         kind: row[7] as string,
+        provenance: row[8] as Provenance,
       },
     }),
     supportedKinds: ["sql", "xss", "logging"],
@@ -100,7 +103,7 @@ export const extensiblePredicateDefinitions: Record<
       method.modeledMethod.input,
       method.modeledMethod.output,
       method.modeledMethod.kind,
-      "manual",
+      method.modeledMethod.provenance,
     ],
     readModeledMethod: (row) => ({
       signature: readRowToMethod(row),
@@ -109,6 +112,7 @@ export const extensiblePredicateDefinitions: Record<
         input: row[6] as string,
         output: row[7] as string,
         kind: row[8] as string,
+        provenance: row[9] as Provenance,
       },
     }),
     supportedKinds: ["taint", "value"],
@@ -124,7 +128,7 @@ export const extensiblePredicateDefinitions: Record<
       method.externalApiUsage.methodName,
       method.externalApiUsage.methodParameters,
       method.modeledMethod.kind,
-      "manual",
+      method.modeledMethod.provenance,
     ],
     readModeledMethod: (row) => ({
       signature: `${row[0]}.${row[1]}#${row[2]}${row[3]}`,
@@ -133,6 +137,7 @@ export const extensiblePredicateDefinitions: Record<
         input: "",
         output: "",
         kind: row[4] as string,
+        provenance: row[5] as Provenance,
       },
     }),
     supportedKinds: ["summary", "source", "sink"],

extensions/ql-vscode/src/stories/data-extensions-editor/DataExtensionsEditor.stories.tsx

@@ -211,30 +211,35 @@ DataExtensionsEditor.args = {
       input: "Argument[0]",
       output: "",
       kind: "jndi-injection",
+      provenance: "df-generated",
     },
     "org.sql2o.Connection#createQuery(String)": {
       type: "summary",
       input: "Argument[this]",
       output: "ReturnValue",
       kind: "taint",
+      provenance: "df-manual",
     },
     "org.sql2o.Sql2o#open()": {
       type: "summary",
       input: "Argument[this]",
       output: "ReturnValue",
       kind: "taint",
+      provenance: "manual",
     },
     "org.sql2o.Query#executeScalar(Class)": {
       type: "neutral",
       input: "",
       output: "",
       kind: "",
+      provenance: "df-generated",
     },
     "org.sql2o.Sql2o#Sql2o(String,String,String)": {
       type: "neutral",
       input: "",
       output: "",
       kind: "",
+      provenance: "df-generated",
     },
   },
 };

extensions/ql-vscode/src/stories/data-extensions-editor/MethodRow.stories.tsx

@@ -50,5 +50,6 @@ MethodRow.args = {
     input: "Argument[this]",
     output: "ReturnValue",
     kind: "taint",
+    provenance: "manual",
   },
 };

extensions/ql-vscode/src/view/data-extensions-editor/MethodRow.tsx

@@ -13,6 +13,7 @@ import { ExternalApiUsage } from "../../data-extensions-editor/external-api-usag
 import {
   ModeledMethod,
   ModeledMethodType,
+  Provenance,
 } from "../../data-extensions-editor/modeled-method";
 import { KindInput } from "./KindInput";
 import { extensiblePredicateDefinitions } from "../../data-extensions-editor/predicates";
@@ -63,13 +64,21 @@ export const MethodRow = ({
     (e: InputEvent) => {
       const target = e.target as HTMLSelectElement;
 
+      let newProvenance: Provenance = "manual";
+      if (modeledMethod?.provenance === "df-generated") {
+        newProvenance = "df-manual";
+      } else if (modeledMethod?.provenance === "ai-generated") {
+        newProvenance = "ai-manual";
+      }
+
       onChange(externalApiUsage, {
         // If there are no arguments, we will default to "Argument[this]"
         input: argumentsList.length === 0 ? "Argument[this]" : "Argument[0]",
         output: "ReturnType",
         kind: "value",
         ...modeledMethod,
         type: target.value as ModeledMethodType,
+        provenance: newProvenance,
       });
     },
     [onChange, externalApiUsage, modeledMethod, argumentsList],

extensions/ql-vscode/test/unit-tests/data-extensions-editor/auto-model.test.ts

@@ -191,12 +191,14 @@ describe("createAutoModelRequest", () => {
       kind: "",
       input: "",
       output: "",
+      provenance: "manual",
     },
     "org.sql2o.Sql2o#Sql2o(String)": {
       type: "sink",
       kind: "jndi-injection",
       input: "Argument[0]",
       output: "",
+      provenance: "manual",
     },
   };
 
@@ -407,18 +409,21 @@ describe("parsePredictedClassifications", () => {
         kind: "sql injection sink",
         input: "Argument[0]",
         output: "",
+        provenance: "ai-generated",
       },
       "org.sql2o.Sql2o#executeScalar(Class)": {
         type: "neutral",
-        kind: "",
+        kind: "summary",
         input: "",
         output: "",
+        provenance: "ai-generated",
       },
       "org.sql2o.Sql2o#Sql2o(String,String,String)": {
         type: "sink",
         kind: "sql injection sink",
         input: "Argument[1]",
         output: "",
+        provenance: "ai-generated",
       },
     });
   });

extensions/ql-vscode/test/unit-tests/data-extensions-editor/yaml.test.ts

@@ -75,6 +75,7 @@ describe("createDataExtensionYaml", () => {
           input: "Argument[0]",
           output: "",
           kind: "sql",
+          provenance: "df-generated",
         },
       },
     );
@@ -89,7 +90,7 @@ describe("createDataExtensionYaml", () => {
       pack: codeql/java-all
       extensible: sinkModel
     data:
-      - ["org.sql2o","Connection",true,"createQuery","(String)","","Argument[0]","sql","manual"]
+      - ["org.sql2o","Connection",true,"createQuery","(String)","","Argument[0]","sql","df-generated"]
 
   - addsTo:
       pack: codeql/java-all
@@ -171,6 +172,7 @@ describe("loadDataExtensionYaml", () => {
         kind: "sql",
         output: "",
         type: "sink",
+        provenance: "manual",
       },
     });
   });