fix: list control suggestions on threats

klarna-incubator · Oct 20, 2023 · e3098a5 · e3098a5
1 parent e438fec
commit e3098a5
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 16 deletions.
diff --git a/app/src/components/model/panels/right/ControlTab.js b/app/src/components/model/panels/right/ControlTab.js
@@ -1,7 +1,7 @@
 import { Box } from "@mui/material";
-import { useSelector } from "react-redux";
 import { useCreateControlMutation } from "../../../../api/gram/controls";
 import { useReadOnly } from "../../../../hooks/useReadOnly";
+import { useModelID } from "../../hooks/useModelID";
 import { useSelectedComponent } from "../../hooks/useSelectedComponent";
 import { useSelectedComponentControls } from "../../hooks/useSelectedComponentControls";
 import { Control } from "./Control";
@@ -11,11 +11,7 @@ export function ControlTab(props) {
   const { scrollToId, selectedId } = props;
   const [createControl] = useCreateControlMutation();
 
-  const { modelId } = useSelector(({ model }) => {
-    return {
-      modelId: model.id,
-    };
-  });
+  const modelId = useModelID();
   const component = useSelectedComponent();
   const controls = useSelectedComponentControls();
 

diff --git a/app/src/components/model/panels/right/Threat.js b/app/src/components/model/panels/right/Threat.js
@@ -21,6 +21,11 @@ import { useModelID } from "../../hooks/useModelID";
 import { EditableSelect } from "./EditableSelect";
 import { EditableTypography } from "./EditableTypography";
 import { MitigationChip } from "./MitigationChip";
+import {
+  useAcceptSuggestionMutation,
+  useListSuggestionsQuery,
+} from "../../../../api/gram/suggestions";
+import { useSelectedComponent } from "../../hooks/useSelectedComponent";
 
 export function Threat({
   threat,
@@ -29,14 +34,29 @@ export function Threat({
   readOnly: propReadOnly,
 }) {
   const modelId = useModelID();
+  const selectedComponent = useSelectedComponent();
   const [deleteThreat] = useDeleteThreatMutation();
   const [updateThreat] = useUpdateThreatMutation();
   const [createControl] = useCreateControlMutation();
   const [createMitigation] = useCreateMitigationMutation();
+  const [acceptSuggestion] = useAcceptSuggestionMutation();
 
   const [title, setTitle] = useState(threat.title);
   const [description, setDescription] = useState(threat.description);
 
+  const partialThreatId = threat.suggestionId
+    ? threat.suggestionId.split("/").splice(1).join("/")
+    : "";
+  const { data: suggestions } = useListSuggestionsQuery(modelId);
+  const controlSuggestions = (
+    suggestions?.controlsMap[selectedComponent.id] || []
+  ).filter(
+    (s) =>
+      partialThreatId &&
+      s.status === "new" &&
+      s.mitigates.find((m) => m.partialThreatId === partialThreatId)
+  );
+
   const controls = useComponentControls(threat.componentId);
   const { data: mitigations } = useListMitigationsQuery({ modelId });
   const threatsMap = mitigations?.threatsMap || {};
@@ -72,11 +92,18 @@ export function Threat({
   }
 
   function onSelectExisting(control) {
-    createMitigation({
-      modelId: threat.modelId,
-      threatId: threat.id,
-      controlId: control.id,
-    });
+    if (control.mitigates) {
+      acceptSuggestion({
+        modelId: modelId,
+        suggestionId: control.id,
+      });
+    } else {
+      createMitigation({
+        modelId: threat.modelId,
+        threatId: threat.id,
+        controlId: control.id,
+      });
+    }
   }
 
   const controlIds = threatsMap[threat.id];
@@ -223,9 +250,12 @@ export function Threat({
         {!readOnly && (
           <EditableSelect
             placeholder="Add Control"
-            options={controls.filter(
-              (c) => !linkedControls.map((l) => l.id).includes(c.id)
-            )}
+            options={[
+              ...controlSuggestions,
+              ...controls.filter(
+                (c) => !linkedControls.map((l) => l.id).includes(c.id)
+              ),
+            ]}
             selectExisting={onSelectExisting}
             createNew={createControlWithMitigation}
           />

diff --git a/plugins/threatlib/src/controls.ts b/plugins/threatlib/src/controls.ts
@@ -43,7 +43,7 @@ const cspHeaders: ThreatLibControlSuggestion = {
   slug: "csp-headers",
   title: "Content Security Policy Headers",
   description: `Configure Content Security Policy headers on your webbapplication to harden against XSS.`,
-  mitigates: ["xss"],
+  mitigates: ["xss", "non-existent"],
 
   condition: (model: Model, component: Component) => {
     return !!component.classes?.find((c) =>
@@ -64,7 +64,9 @@ export function mapControls(
     .filter((c) => c.condition(model, component))
     .map((c) => ({
       ...c,
-      mitigates: c.mitigates.map((m) => ({ partialThreatId: m })),
+      mitigates: c.mitigates.map((m) => ({
+        partialThreatId: m.includes("/") ? m : `threatlib/threat/${m}`,
+      })),
       componentId: component.id,
     }));
 }