Automative threat generation for LLM-agent based systems

.gitignore

@@ -118,6 +118,7 @@
 !td.vue/**/*.scss
 !td.vue/public/
 !td.vue/public/*
+!td.vue/public/**
 !td.vue/src/
 !td.vue/src/*.js
 !td.vue/src/*.vue
@@ -166,6 +167,9 @@
 !td.vue/src/service/threats/models/*.js
 !td.vue/src/service/threats/oats/
 !td.vue/src/service/threats/oats/*.js
+!td.vue/src/service/threats/analyzer/
+!td.vue/src/service/threats/analyzer/*.js
+!td.vue/src/service/threats/analyzer/*.json
 !td.vue/src/service/x6/
 !td.vue/src/service/x6/*.js
 !td.vue/src/service/x6/graph/

td.vue/src/components/GraphButtons.vue

@@ -1,5 +1,12 @@
 <template>
     <b-btn-group>
+
+        <td-form-button
+            :onBtnClick="generateThreats"
+            icon="magic-wand-sparkles"
+            :title="$t('threatmodel.buttons.generateThreats')"
+            text="" />
+
         <td-form-button
             :onBtnClick="deleteSelected"
             icon="trash"
@@ -73,14 +80,25 @@
 import { mapState } from 'vuex';
 
 import TdFormButton from '@/components/FormButton.vue';
+import { completeGraphAgentThreats } from '@/service/threats/analyzer/index.js';
+
 
 export default {
     name: 'TdGraphButtons',
     components: {
         TdFormButton
     },
     computed: mapState({
+        cellRef: (state) => state.cell.ref,
+        threats: (state) => state.cell.threats,
         diagram: (state) => state.threatmodel.selectedDiagram,
+        threatTop: (state) => state.threatmodel.data.detail.threatTop,
+        disableNewThreat: function (state) {
+            if (!state.cell?.ref?.data) {
+                return true;
+            }
+            return state.cell.ref.data.outOfScope || state.cell.ref.data.isTrustBoundary || state.cell.ref.data.type === 'tm.Text';
+        }
     }),
     data() {
         return {
@@ -93,6 +111,9 @@ export default {
         }
     },
     methods: {
+        generateThreats() {
+            completeGraphAgentThreats(this.graph, this.$store);
+        },
         save() {
             this.$emit('saved');
         },

td.vue/src/components/GraphProperties.vue

@@ -242,6 +242,22 @@
             >
           </b-form-group>
         </b-col>
+
+        <!-- Auto properties -->
+        <b-col 
+            v-if="cellRef.data.type in properties"
+            v-for="prop in properties[cellRef.data.type]" 
+            :key="`${cellRef.data.type}_${prop.key}`">
+          <b-form-group label-cols="auto" :id="`${prop.key}-group`">
+            <component
+              :is="propertyTypeComponents[prop.type]"
+              :id="prop.key"
+              v-model="cellRef.data[prop.key]"
+              @change="onChangeProperties()"
+              >{{ $t(`threatmodel.properties.${prop.key}`) }}</component>
+          </b-form-group>
+        </b-col>
+
       </b-form-row>
     </b-form>
   </div>
@@ -259,6 +275,23 @@ import dataChanged from '@/service/x6/graph/data-changed.js';
 
 export default {
     name: 'TdGraphProperties',
+    data: () => ({
+        properties: {
+            'tm.Agent': [
+                { key: 'pQueryRewriting', type: 'bool' },
+                { key: 'pPromptTemplate', type: 'bool' },
+                { key: 'pAttachInstruction', type: 'bool' },
+                { key: 'pOutputSafetyFilter', type: 'bool' },
+                { key: 'pModifiesMemory', type: 'bool' },
+            ],
+            'tm.Flow': [
+                { key: 'pRawUserData', type: 'bool' },
+            ]
+        },
+        propertyTypeComponents: {
+            bool: 'b-form-checkbox'
+        }
+    }),
     computed: mapState({
         cellRef: (state) => state.cell.ref
     }),

td.vue/src/components/Navbar.vue

@@ -22,6 +22,7 @@
             v-b-tooltip.hover :title="$t('nav.logOut')"
           ></font-awesome-icon>
         </b-nav-item>
+
         <b-nav-item
           href="https://www.threatdragon.com/docs/"
           target="_blank"
@@ -136,7 +137,8 @@ export default {
                     throw error;
                 }
             });
-        }
+        },
+
     }
 };
 </script>

td.vue/src/i18n/en.js

@@ -181,10 +181,17 @@ const eng = {
             protocol: 'Protocol',
             publicNetwork: 'Public Network',
             storesCredentials: 'Stores Credentials',
-            storesInventory: 'Stores Inventory'
+            storesInventory: 'Stores Inventory',
+            pQueryRewriting: 'Query Rewriting',
+            pPromptTemplate: 'Prompt Template',
+            pAttachInstruction: 'Attach Instruction',
+            pOutputSafetyFilter: 'Output Safety Filter',
+            pModifiesMemory: 'Modifies Memory',
+            pRawUserData: 'Raw User Data',
         },
         buttons: {
             delete: 'Delete selected',
+            generateThreats: 'Generate threats',
             redo: 'Redo edit',
             shortcuts: 'Keyboard shortcuts',
             toggleGrid: 'Toggle grid',
@@ -241,6 +248,8 @@ const eng = {
         },
         shapes: {
             actor: 'Actor',
+            agent: 'Agent',
+            tool: 'Tool',
             flow: 'Data Flow',
             flowStencil: 'Data Flow',
             process: 'Process',

td.vue/src/i18n/ru.js

@@ -180,6 +180,7 @@ const rus = {
         },
         buttons: {
             delete: 'Delete selected',
+            generateThreats: 'Generate threats',
             redo: 'Redo edit',
             shortcuts: 'Keyboard shortcuts',
             toggleGrid: 'Toggle grid',
@@ -236,6 +237,8 @@ const rus = {
         },
         shapes: {
             actor: 'Actor',
+            agent: 'Agent',
+            tool: 'Tool',
             flow: 'Data Flow',
             flowStencil: 'Data Flow',
             process: 'Process',

td.vue/src/plugins/fontawesome-vue.js

@@ -31,7 +31,8 @@ import {
     faPrint,
     faProjectDiagram,
     faDiagramProject,
-    faLock
+    faLock,
+    faMagicWandSparkles
 } from '@fortawesome/free-solid-svg-icons';
 
 import {faBitbucket, faGithub, faGitlab, faVuejs, faGoogle, faGoogleDrive} from '@fortawesome/free-brands-svg-icons';
@@ -70,7 +71,8 @@ library.add(
     faDiagramProject,
     faGoogle,
     faGoogleDrive,
-    faLock
+    faLock,
+    faMagicWandSparkles
 );
 
 Vue.component('font-awesome-icon', FontAwesomeIcon);

td.vue/src/service/entity/default-properties.js

@@ -1,28 +1,3 @@
-const actor = {
-    type: 'tm.Actor',
-    name: 'Actor',
-    description: '',
-    outOfScope: false,
-    reasonOutOfScope: '',
-    hasOpenThreats: false,
-    providesAuthentication: false,
-    threats: []
-};
-
-const boundary = {
-    type: 'tm.Boundary',
-    name: 'Trust Boundary',
-    description: '',
-    isTrustBoundary: true
-};
-
-const boundaryBox = {
-    type: 'tm.BoundaryBox',
-    name: 'Trust Boundary',
-    description: '',
-    isTrustBoundary: true
-};
-
 const flow = {
     type: 'tm.Flow',
     name: 'Data Flow',
@@ -34,7 +9,23 @@ const flow = {
     isEncrypted: false,
     isPublicNetwork: false,
     protocol: '',
-    threats: []
+    pRawUserData: false,
+    threats: [],
+    "isFunctionCall": false,
+    "isHighFrequency": false,
+    "resultsAreValidated": true
+};
+
+const actor = {
+    type: 'tm.Actor',
+    name: 'Actor',
+    description: '',
+    outOfScope: false,
+    reasonOutOfScope: '',
+    hasOpenThreats: false,
+    providesAuthentication: false,
+    threats: [],
+    isTrusted: false
 };
 
 const tmProcess = {
@@ -48,7 +39,11 @@ const tmProcess = {
     handlesGoodsOrServices: false,
     isWebApplication: false,
     privilegeLevel: '',
-    threats: []
+    threats: [],
+    isTrusted: false,
+    "hasShortTermMemory": false,
+    "hasLongTermMemory": false,
+    "executesGeneratedCode": false
 };
 
 const store = {
@@ -63,16 +58,73 @@ const store = {
     isSigned: false,
     storesCredentials: false,
     storesInventory: false,
-    threats: []
+    threats: [],
+    isTrusted: false,
+    "isSharedContext": false
+};
+
+const agent = {
+    type: 'tm.Agent',
+    name: 'Agent',
+    description: '',
+    outOfScope: false,
+    reasonOutOfScope: '',
+    hasOpenThreats: false,
+    pQueryRewriting: false,
+    pPromptTemplate: false,
+    pAttachInstruction: false,
+    pOutputSafetyFilter: false,
+    pModifiesMemory: false,
+    threats: [],
+    isTrusted: false,
+    "hasShortTermMemory": false,
+    "hasLongTermMemory": false,
+    "trustScore": 50,
+    "isRogue": false
 };
 
+const tool = {
+    type: 'tm.Tool',
+    name: 'Tool',
+    description: '',
+    outOfScope: false,
+    reasonOutOfScope: '',
+    hasOpenThreats: false,
+    pRequiresElevatedPrivileges: false,
+    threats: [],
+    isTrusted: false,
+    "isExternal": false,
+    "isDangerous": false,
+    "requiresElevatedPrivileges": false
+};
+
+
+
+
+
 const text = {
     type: 'tm.Text',
     name: 'Descriptive text'
 };
 
+const boundary = {
+    type: 'tm.Boundary',
+    name: 'Trust Boundary',
+    description: '',
+    isTrustBoundary: true
+};
+
+const boundaryBox = {
+    type: 'tm.BoundaryBox',
+    name: 'Trust Boundary',
+    description: '',
+    isTrustBoundary: true
+};
+
 const propsByType = {
     'tm.Actor': actor,
+    'tm.Agent': agent,
+    'tm.Tool': tool,
     'tm.Boundary': boundary,
     'tm.BoundaryBox': boundaryBox,
     'tm.Flow': flow,
@@ -91,6 +143,7 @@ const getByType = (type) => {
 
 export default {
     actor,
+    agent,
     boundary,
     boundaryBox,
     flow,