cds-snc · maxneuvians · Mar 30, 2024 · Apr 3, 2024 · Apr 12, 2024 · Apr 12, 2024
diff --git a/package.json b/package.json
@@ -67,6 +67,7 @@
     "eslint-plugin-header": "^3.1.1",
     "indefinite": "*",
     "lodash": "^4.17.21",
+    "openai": "^4.30.0",
     "react": "*",
     "react-dom": "*",
     "react-draggable": "^4.4.6",

diff --git a/public/index.html b/public/index.html
@@ -4,7 +4,7 @@
 <head>
   <meta charset="utf-8" />
   <meta http-equiv='Content-Security-Policy'
-    content="default-src 'self' ; img-src 'self' * data:; font-src 'self' data:; script-src 'self'; style-src 'self' 'unsafe-inline'; media-src 'none' ; frame-src 'none';" />
+    content="default-src 'self' https://api.openai.com/v1/chat/completions; img-src 'self' * data:; font-src 'self' data:; script-src 'self'; style-src 'self' 'unsafe-inline'; media-src 'none' ; frame-src 'none';" />
   <meta name="viewport" content="width=device-width, initial-scale=1" />
   <meta name="theme-color" content="#000000" />
   <meta name="description" content="threat-composer" />

diff --git a/src/components/application/ApplicationInfo/index.tsx b/src/components/application/ApplicationInfo/index.tsx
@@ -35,6 +35,7 @@ const ApplicationInfo: FC<EditableComponentBaseProps> = ({
   const [editMode, setEditMode] = useState(!applicationInfo.name && !applicationInfo.description );
   const [content, setContent] = useState('');
   const [name, setName] = useState('');
+  const [openAIKey, setOpenAIKey] = useState('');
   const [securityCategory, setSecurityCategory] = useState('CCCS Medium');
   const [checkedIaaS, setCheckedIaaS] = useState(false);
   const [checkedPaaS, setCheckedPaaS] = useState(false);
@@ -54,6 +55,7 @@ const ApplicationInfo: FC<EditableComponentBaseProps> = ({
       ...prev,
       description: content,
       name,
+      openAIKey,
       securityCategory: securityCategory,
       useIaaS: checkedIaaS,
       usePaaS: checkedPaaS,
@@ -67,11 +69,12 @@ const ApplicationInfo: FC<EditableComponentBaseProps> = ({
     setEditMode(false);
   }, [checkedApplication, checkedCompute, checkedData,
     checkedIaaS, checkedNetwork, checkedPaaS, checkedSaaS,
-    checkedStorage, content, name, securityCategory, setApplicationInfo]);
+    checkedStorage, content, name, openAIKey, securityCategory, setApplicationInfo]);
 
   const handleEdit = useCallback(() => {
     setContent(applicationInfo.description || '');
     setName(applicationInfo.name || '');
+    setOpenAIKey(applicationInfo.openAIKey || '');
     setSecurityCategory(applicationInfo.securityCategory || 'CCCS Medium');
     setCheckedIaaS(applicationInfo.useIaaS || false);
     setCheckedPaaS(applicationInfo.usePaaS || false);
@@ -201,6 +204,18 @@ const ApplicationInfo: FC<EditableComponentBaseProps> = ({
           Network
           </Checkbox>
         </FormField>
+        <FormField
+          label="OpenAI Key"
+        >
+          <Input
+            value={openAIKey}
+            onChange={event =>
+              setOpenAIKey(event.detail.value)
+            }
+            validateData={ApplicationInfoSchema.shape.openAIKey.safeParse}
+            placeholder='Enter OpenAI API key'
+          />
+        </FormField>
       </SpaceBetween>) :
       (
         <SpaceBetween direction='vertical' size='l'>

diff --git a/src/components/generic/Flow/Nodes/ProcessNode.tsx b/src/components/generic/Flow/Nodes/ProcessNode.tsx
@@ -26,15 +26,13 @@ export default memo(({ data, selected }: { data: any; selected: boolean }) => {
         id="top-source"
         position={Position.Top}
         style={{ background: '#555' }}
-        onConnect={(params) => console.log('handle onConnect', params)}
         isConnectable={true}
       />
       <Handle
         type="source"
         id="left-source"
         position={Position.Left}
         style={{ background: '#555' }}
-        onConnect={(params) => console.log('handle onConnect', params)}
         isConnectable={true}
       />
       <ProcessStyle selected={selected}>

diff --git a/src/components/generic/Flow/Threats/ThreatList.tsx b/src/components/generic/Flow/Threats/ThreatList.tsx
@@ -26,7 +26,8 @@ import {
   PropertyFilter,
 } from '@cloudscape-design/components';
 import { columnDefinitions, getMatchesCountText, paginationLabels, collectionPreferencesProps, filteringConstants, filteringProperties } from './table-config';
-import { useThreatsContext } from '../../../../contexts/ThreatsContext';
+import { useThreatsContext, useApplicationInfoContext } from '../../../../contexts';
+import OpenAIModal from '../../OpenAIModal';
 
 function EmptyState({ title, subtitle, action }) {
   return (
@@ -42,9 +43,12 @@ function EmptyState({ title, subtitle, action }) {
   );
 }
 
-export default memo(({ threats, component, changeHandler }: { threats: any; component: any; changeHandler: any } ) => {
+export default memo(({ threats, component, changeHandler, flow }: { threats: any; component: any; changeHandler: any; flow: any }) => {
 
   const [data, setData] = useState((component && component.data) || {});
+  const [AIModalVisable, setAIModalVisable] = useState(false);
+
+  const { applicationInfo } = useApplicationInfoContext();
 
   useEffect(() => {
     setData((component && component.data) || {});
@@ -97,7 +101,7 @@ export default memo(({ threats, component, changeHandler }: { threats: any; comp
             subtitle=""
             action={
               <Button onClick={() => actions.setFiltering('')}>
-              Clear filter
+                Clear filter
               </Button>
             }
           />
@@ -117,48 +121,54 @@ export default memo(({ threats, component, changeHandler }: { threats: any; comp
   }
 
   return (
-    <Table
-      {...collectionProps}
-      selectionType="multi"
-      header={
-        <Header
-          counter={
-            data.threats?.length
-              ? `(${data.threats.length}/${threats.length})`
-              : `(${threats.length})`
-          }
-        >
-          Threats
-        </Header>
-      }
-      columnDefinitions={columnDefinitions}
-      visibleColumns={preferences.visibleContent}
-      items={items}
-      trackBy="id"
-      selectedItems={data.threats}
-      onSelectionChange={(e) => updateData('threats', e.detail.selectedItems, [])}
-      stickyHeader
-      resizableColumns
-      wrapLines
-      stripedRows
-      pagination={
-        <Pagination {...paginationProps} ariaLabels={paginationLabels} />
-      }
-      preferences={
-        <CollectionPreferences
-          {...collectionPreferencesProps}
-          preferences={preferences}
-          onConfirm={({ detail }) => setPreferences(detail)}
-        />
-      }
-      filter={
-        <PropertyFilter
-          {...propertyFilterProps}
-          countText={getMatchesCountText(filteredItemsCount)}
-          expandToViewport={true}
-          i18nStrings={filteringConstants}
-        />
-      }
-    />
+    <>
+      <Table
+        {...collectionProps}
+        selectionType="multi"
+        header={
+          <Header
+            actions={
+              (applicationInfo?.openAIKey ? <Button variant="primary" onClick={() => setAIModalVisable(true)}>Generate</Button> : null)
+            }
+            counter={
+              data.threats?.length
+                ? `(${data.threats.length}/${threats.length})`
+                : `(${threats.length})`
+            }
+          >
+            Threats
+          </Header>
+        }
+        columnDefinitions={columnDefinitions}
+        visibleColumns={preferences.visibleContent}
+        items={items}
+        trackBy="id"
+        selectedItems={data.threats}
+        onSelectionChange={(e) => updateData('threats', e.detail.selectedItems, [])}
+        stickyHeader
+        resizableColumns
+        wrapLines
+        stripedRows
+        pagination={
+          <Pagination {...paginationProps} ariaLabels={paginationLabels} />
+        }
+        preferences={
+          <CollectionPreferences
+            {...collectionPreferencesProps}
+            preferences={preferences}
+            onConfirm={({ detail }) => setPreferences(detail)}
+          />
+        }
+        filter={
+          <PropertyFilter
+            {...propertyFilterProps}
+            countText={getMatchesCountText(filteredItemsCount)}
+            expandToViewport={true}
+            i18nStrings={filteringConstants}
+          />
+        }
+      />
+      <OpenAIModal visible={AIModalVisable} setVisible={setAIModalVisable} data={component} apiKey={applicationInfo?.openAIKey} flow={flow} />
+    </>
   );
 });
diff --git a/src/components/generic/Flow/index.tsx b/src/components/generic/Flow/index.tsx
@@ -22,6 +22,8 @@ import SpaceBetween from '@cloudscape-design/components/space-between';
 import { v4 } from 'uuid';
 import 'reactflow/dist/style.css';
 
+import { LOCAL_STORAGE_KEY_DATA_FLOW_DIAGRAM } from '../../../configs/localStorageKeys';
+
 import ActorNode from './Nodes/ActorNode';
 import DatastoreNode from './Nodes/DatastoreNode';
 import ProcessNode from './Nodes/ProcessNode';
@@ -37,7 +39,6 @@ import ZIndexChanger from './Nodes/ZIndexChanger';
 import ThreatList from './Threats/ThreatList';
 
 import { useThreatsContext } from '../../../contexts';
-import { useWorkspacesContext } from '../../../contexts/WorkspacesContext';
 
 const edgeTypes = {
   biDirectional: BiDirectionalEdge,
@@ -63,9 +64,6 @@ namespace s {
 }
 
 function Flow() {
-  const { currentWorkspace } = useWorkspacesContext();
-  const flowKey = `dataflow-diagram-${currentWorkspace?.id}`;
-
   const { zoomTo, getZoom, setViewport } = useReactFlow();
 
   // Save and restore state
@@ -75,15 +73,15 @@ function Flow() {
   const onSave = useCallback(() => {
     if (rfInstance) {
       const flow = rfInstance.toObject();
-      localStorage.setItem(flowKey, JSON.stringify(flow));
+      localStorage.setItem(LOCAL_STORAGE_KEY_DATA_FLOW_DIAGRAM, JSON.stringify(flow));
       setSaveState(true);
     }
-  }, [rfInstance, flowKey]);
+  }, [rfInstance]);
 
   const onInit = async (instance) => {
     setRfInstance(instance);
     const restoreFlow = async () => {
-      const flow = JSON.parse(localStorage.getItem(flowKey) as string);
+      const flow = JSON.parse(localStorage.getItem(LOCAL_STORAGE_KEY_DATA_FLOW_DIAGRAM) as string);
 
       if (flow) {
         const { x = 0, y = 0, zoom = 1 } = flow.viewport;
@@ -275,7 +273,8 @@ function Flow() {
       <ThreatList
         threats={threatList}
         component={selectedComponent}
-        changeHandler={setNodeDataValue} />
+        changeHandler={setNodeDataValue}
+        flow={rfInstance?.toObject()}/>
     </SpaceBetween>
   );
 }

diff --git a/src/components/generic/OpenAIModal/index.tsx b/src/components/generic/OpenAIModal/index.tsx
@@ -0,0 +1,111 @@
+import { memo, useEffect, useState } from 'react';
+import {
+  Modal,
+  Header,
+  Alert,
+  Box,
+  Table,
+  SpaceBetween,
+  Button,
+} from '@cloudscape-design/components';
+
+import OpenAI from 'openai';
+import { v4 } from 'uuid';
+
+import { useThreatsContext } from '../../../contexts/ThreatsContext/context';
+
+const prompt = () => {
+  return `You are an assitant helping to generate threat statements based on the STRIDE threat modeling framework. Threat statements are structured in a using a specific grammar. The grammar we are using is as follows: [threat source] [prerequisites] can [threat action] which leads to [threat impact], negatively impacting [impacted assets]. We are also using a data flow diagram to help generate these statements. I will provide you with information about the impacted asset in a json format, which is described either as an actor, a process, a datastore, or an edge linking any of those elements. I will also provide you with additional information about the impacted asset such as data features, technology features, security features, and tags. Based on this information and the general context of the diagram, you will generate a threat statement. Please ensure the statement follows the grammar provided. You do not need to include [] in the statement. Ensure that the threat statement you return makes logical sense in the context of threat modeling. Please always return one threat statement in a JSON array of objects that follow this structure: {threatSource: '', prerequisites: '', threatAction: '', threatImpact: '', impactedGoal: [], impactedAssets: [], statement: '', displayedStatement: []}.  For example, assume you generate a threat statement as follows:
+
+  A threat actor with access to logs can access sensitive data that was not excluded or redacted, which leads to the ability to modify data, resulting in reduced confidentiality of file upload, file listing, upload status, uploaded files and authentication tokens
+
+  You would return:
+
+  [{"statement":"An external threat actor who is authenticated and authorized can flood the legitimate user's internet link, which leads to unnecessary and excessive costs, resulting in reduced availability of billing","threatSource":"external threat actor","prerequisites":"who is authenticated and authorized","displayedStatement":["An ",{"type":"span","content":"external threat actor","tooltip":"threat source"}," ",{"type":"span","content":"who is authenticated and authorized","tooltip":"prerequisites"}," can ",{"type":"b","content":"flood the legitimate user's internet link","tooltip":"threat action"},", which leads to ",{"type":"span","content":"unnecessary and excessive costs","tooltip":"threat impact"},", resulting in reduced ",{"type":"span","content":"availability","tooltip":"impacted goal"}," of ",{"type":"span","content":"billing","tooltip":"impacted assets"}],"threatAction":"flood the legitimate user's internet link","threatImpact":"unnecessary and excessive costs","impactedAssets":["billing"],"impactedGoal":["availability"]}]
+
+  ONLY RETURN VALID JSON THAT IS PARSEABLE BY A JAVASCRIPT JSON PARSER. DO NOT INCLUDE \`\`\`json TAGS IN YOUR RESPONSE. DO NOT INCLUDE ANYTHING OTHER THAN THE JSON ARRAY OF OBJECTS.
+  `;
+};
+
+export default memo(({ visible, setVisible, data, flow, apiKey }: { visible: boolean; setVisible: any; data: any; flow: any; apiKey: any }) => {
+
+  const [loading, setLoading] = useState(false);
+  const [threats, setThreats] = useState([]);
+  const [error, setError] = useState(null);
+  const [selectedItems, setSelectedItems] = useState<{ statement: string }[]>([]);
+
+  const { saveStatement } = useThreatsContext();
+
+  useEffect(() => {
+    if (visible) {
+      const openai = new OpenAI({ apiKey: apiKey, dangerouslyAllowBrowser: true });
+      setError(null);
+      setLoading(true);
+      openai.chat.completions.create({
+        messages: [
+          { role: 'system', content: prompt() },
+          { role: 'assistant', content: 'Generate a threat statement for the impacted asset.' },
+          { role: 'user', content: `Here is the information about the impacted asset in json format: ${JSON.stringify(data)} Here is the data flow diagram in json format: ${JSON.stringify(flow)}` },
+        ],
+        model: 'gpt-4o',
+      }).then((response) => {
+        let content = response.choices[0].message.content || '[]';
+        setThreats(JSON.parse(content).map((item) => ({ id: v4(), ...item })) || []);
+        setLoading(false);
+        return response;
+      }).catch((err) => {
+        setError(err.message);
+        setLoading(false);
+      });
+    };
+  }, [visible, apiKey, data, flow]);
+
+  const saveSuggestedThreats = () => {
+    selectedItems.forEach((item) => {
+      let s = {
+        id: v4(),
+        numericId: -1,
+        ...item,
+      };
+      saveStatement(s);
+    });
+    setVisible(false);
+  };
+
+
+  return <Modal
+    header={<Header>AI generated threat statements</Header>}
+    visible={visible}
+    onDismiss={() => setVisible(false)}
+    size='large'
+    footer={
+      <Box float="right">
+        <SpaceBetween direction="horizontal" size="xs">
+          <Button variant="link" onClick={() => setVisible(false)}>Cancel</Button>
+          <Button variant="primary" onClick={() => saveSuggestedThreats()}>Add selected statements</Button>
+        </SpaceBetween>
+      </Box>}
+  >
+    {(error ? <Alert
+      statusIconAriaLabel='Error'
+      type='error'
+      header='There was an error calling the API'>
+      {error}
+    </Alert> :
+      <Table
+        columnDefinitions={[{ id: 'statement', header: 'Statement', cell: (item: { id: number; statement: string }) => item.statement || '-' }]}
+        enableKeyboardNavigation
+        items={threats}
+        loading={loading}
+        loadingText='Loading statements...'
+        sortingDisabled
+        trackBy="id"
+        variant='embedded'
+        selectedItems={selectedItems}
+        onSelectionChange={({ detail }) => setSelectedItems(detail.selectedItems)}
+        selectionType="multi"
+        empty={<Box margin={{ vertical: 'xs' }} textAlign='center' color='inherit'></Box>}
+        wrapLines />
+    )}
+  </Modal>;
+});
diff --git a/src/configs/localStorageKeys.ts b/src/configs/localStorageKeys.ts
@@ -35,5 +35,7 @@ export const LOCAL_STORAGE_KEY_ARCHIECTURE_INFO = 'ThreatStatementGenerator.Arch
 export const LOCAL_STORAGE_KEY_DATAFLOW_INFO = 'ThreatStatementGenerator.DataflowInfo';
 export const LOCAL_STORAGE_KEY_DIAGRAM_INFO = 'ThreatStatementGenerator.DiagramInfo';
 
+export const LOCAL_STORAGE_KEY_DATA_FLOW_DIAGRAM = 'ThreatStatementGenerator.DataFlowDiagram';
+
 export const LOCAL_STORAGE_KEY_WORKSPACE_LIST_MIGRATION = 'ThreatStatementGenerator.workspaceListMigration';
 export const LOCAL_STORAGE_KEY_THREATS_LIST_MIGRATION = 'ThreatStatementGenerator.threatListMigration';
diff --git a/src/contexts/ApplicationContext/context.ts b/src/contexts/ApplicationContext/context.ts
@@ -27,6 +27,7 @@ export interface ApplicationInfoContextApi {
 const initialState: ApplicationInfoContextApi = {
   applicationInfo: {
     description: '',
+    openAIKey: '',
     securityCategory: 'CCCS Medium',
     useIaaS: false,
     usePaaS: false,