fix: Fix an issue where a tc file cannot be reimported if it contains…

… threat pack data (#141)
awslabs · Sep 24, 2024 · 180fa1e · 180fa1e
1 parent cb50ccb
commit 180fa1e
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 8 deletions.
diff --git a/packages/threat-composer/src/components/threats/MitigationCandidates/index.tsx b/packages/threat-composer/src/components/threats/MitigationCandidates/index.tsx
@@ -116,7 +116,10 @@ const MitigationCandidates: FC<MitigationCandidatesProp> = ({
     variant='embedded'
     items={items || []}
     selectedItems={selectedItems}
-    onSelectionChange={({ detail }) => setSelectedItems([...detail.selectedItems])}
+    onSelectionChange={({ detail }) => setSelectedItems([...detail.selectedItems.map(x => {
+      const { comments, ...data } = x;
+      return data;
+    })])}
     isItemDisabled={(item) => linkedMitigationsFromThreakpack.includes(item.id)}
   />);
 };

diff --git a/packages/threat-composer/src/components/workspaces/MitigationPack/index.tsx b/packages/threat-composer/src/components/workspaces/MitigationPack/index.tsx
@@ -37,10 +37,15 @@ const MitigationPack: FC<MitigationPackProp> = ({
     return mitigationPacks.find(x => x.id === mitigationPackId);
   }, []);
 
-  const [selectedItems, setSelectedItems] = useState<Mitigation[]>([]);
+  const [selectedItems, setSelectedItems] = useState<(Mitigation & {
+    comments?: string;
+  })[]>([]);
 
   const handleAddToWorkspace = useCallback(async () => {
-    await addMitigations(mitigationPackId, selectedItems);
+    await addMitigations(mitigationPackId, selectedItems.map(x => {
+      const { comments, ...data } = x;
+      return data;
+    }));
     setSelectedItems([]);
   }, [mitigationPackId, selectedItems]);
 

diff --git a/packages/threat-composer/src/configs/metadata.ts b/packages/threat-composer/src/configs/metadata.ts
@@ -24,4 +24,26 @@ export const METADATA_SOURCE_MITIGATION_PACK = 'mitigationPack';
 export const METADATA_KEY_SOURCE_MITIGATION_PACK = 'mitigationPackId';
 export const METADATA_KEY_SOURCE_MITIGATION_PACK_MITIGATION = 'mitigationPackMitigationId';
 
-export const METADATA_KEY_DESCRIPTION = 'Description';
+export const METADATA_KEY_DESCRIPTION = 'Description';
+
+export const METADATA_KEY_COMMENTS = 'Comments';
+export const METADATA_KEY_STRIDE = 'STRIDE';
+export const METADATA_KEY_PRIORITY = 'Priority';
+
+export const METADATA_KEY_PREFIX_CUSTOM = 'custom:';
+
+export const ALLOW_METADATA_TAGS = [
+  METADATA_KEY_COMMENTS,
+  METADATA_KEY_PRIORITY,
+  METADATA_KEY_STRIDE,
+
+  METADATA_KEY_SOURCE,
+  METADATA_KEY_SOURCE_THREAT_PACK,
+  METADATA_KEY_SOURCE_THREAT_PACK_THREAT,
+  METADATA_KEY_SOURCE_THREAT_PACK_MITIGATION_CANDIDATE,
+
+  METADATA_KEY_SOURCE_MITIGATION_PACK,
+  METADATA_KEY_SOURCE_MITIGATION_PACK_MITIGATION,
+
+  METADATA_KEY_DESCRIPTION,
+];
diff --git a/packages/threat-composer/src/customTypes/entities.ts b/packages/threat-composer/src/customTypes/entities.ts
@@ -24,6 +24,11 @@ import {
   REGEX_CONTENT_IMAGE_BASE64,
   IMAGE_BASE64_MAX_LENGTH,
   IMAGE_URL_MAX_LENGTH,
+  METADATA_KEY_COMMENTS,
+  METADATA_KEY_STRIDE,
+  METADATA_KEY_PRIORITY,
+  ALLOW_METADATA_TAGS,
+  METADATA_KEY_PREFIX_CUSTOM,
 } from '../configs';
 import STRIDE from '../data/stride';
 
@@ -35,19 +40,23 @@ export const MetadataSchema = z.object({
   key: z.string().max(SINGLE_FIELD_INPUT_SMALL_MAX_LENGTH),
   value: z.union([z.string(), z.array(z.string())]),
 }).strict().refine((data) => {
-  if (data.key === 'Comments') {
+  if (!ALLOW_METADATA_TAGS.includes(data.key) && !data.key.startsWith(METADATA_KEY_PREFIX_CUSTOM)) {
+    return false;
+  }
+
+  if (data.key === METADATA_KEY_COMMENTS) {
     return MetadataCommentSchema.safeParse(data.value).success;
   }
 
-  if (data.key === 'STRIDE') {
+  if (data.key === METADATA_KEY_STRIDE) {
     return Array.isArray(data.value) && data.value.every(v => STRIDE.map(s => s.value).includes(v));
   }
 
-  if (data.key === 'Priority') {
+  if (data.key === METADATA_KEY_PRIORITY) {
     return typeof data.value === 'string' && LEVEL_SELECTOR_OPTIONS.map(o => o.value).includes(data.value);
   }
 
-  return false;
+  return true;
 }, (data) => ({
   message: `Invalid key ${data.key} with value ${JSON.stringify(data.value)}`,
   path: [data.key],