LLM-powered threat modeling
Techniques for implementing an AI threat modeling tool are disclosed. A static analysis tool is used to extract a candidate code snippet from a code repository. The candidate code snippet is identified as potentially being a security relevant code element. The static analysis tool generates additional context associated with the candidate code snippet. An LLM prompt is generated. This prompt is structured to include the candidate code snippet, the context, and a directive to assign a classification to the candidate code snippet. The classification includes a source classification, a sink classification, a sanitizer classification, or a flow step classification. The LLM operates on the prompt to generate output comprising a specific classification for the candidate code snippet. The output is formatted into a data extension file that is consumable by the static analysis tool.
1 . A method for implementing an artificial intelligence (Al) threat modeling tool, said method being performed by a computing service, wherein the computing service is one of a cloud service operating in a cloud environment, a local service operating on a local device, or a hybrid service that includes a cloud component operating in the cloud environment and a local component operating on the local device, said method comprising:
using a static analysis tool to extract a candidate code snippet from a code repository, wherein the candidate code snippet is identified by the static analysis tool as potentially being a security relevant code element;
using the static analysis tool to generate additional context associated with the candidate code snippet;
generating a large language model (LLM) prompt, which is structured to include the candidate code snippet, the context, and a directive to assign a classification to the candidate code snippet, said classification including a source classification, a sink classification, a sanitizer classification, or a flow step classification, and wherein the directive further includes a directive to determine a type for the classification;
triggering the LLM to operate on the LLM prompt, wherein, as a result of said operating, the LLM generates output comprising a specific classification for the candidate code snippet and a specific type for the specific classification;
formatting the output of the LLM into one or more data extension files that are consumable by the static analysis tool; and
including the one or more data extension files in a corpus of data extension files that are consumable by the static analysis tool.
2 . The method of claim 1 , wherein the context includes one or more few-shot examples.
3 . The method of claim 2 , wherein the one or more few-shot examples include at least one of: a positive example of a sink, a positive example of a source, a positive example of a sanitizer, or a positive example of a flow step.
4 . The method of claim 2 , wherein the one or more few-shot examples include at least one of: a negative example of a sink, a negative example of a source, a negative example of a sanitizer, or a negative example of a flow step.
5 . The method of claim 1 , wherein the context further includes documentation associated with the candidate code snippet.
6 . The method of claim 1 , wherein the context further includes content extracted from a selected number of code lines preceding a code line comprising the candidate code snippet.
7 . The method of claim 1 , wherein the context further includes content extracted from a selected number of code lines succeeding a code line comprising the candidate code snippet.
8 . The method of claim 1 , wherein formatting the output of the LLM into the data extension file includes parsing the output and organizing the parsed output into a format that is consumable by the static analysis tool.
9 . The method of claim 1 , wherein the candidate code snippet is one of a plurality of candidate code snippets that are identified and extracted by the static analysis tool, and wherein a pre-filtering operation is performed on the plurality of candidate code snippets using a defined set of parameters to reduce a number of candidate code snippets included in the plurality of candidate code snippets.
10 . The method of claim 1 , wherein the prompt includes details regarding a particular classification provided to a few-shot example that is also included in the prompt.
11 . A computer system comprising:
a processor system; and
a storage system that includes instructions that are executable by the processor system to cause the computer system to:
use a static analysis tool to extract a candidate code snippet from a code repository, wherein the candidate code snippet is identified by the static analysis tool as potentially being a security relevant code element;
use the static analysis tool to generate additional context associated with the candidate code snippet;
generate a large language model (LLM) prompt, which is structured to include the candidate code snippet, the context, and a directive to assign a classification to the candidate code snippet, said classification including a source classification, a sink classification, a sanitizer classification, or a flow step classification, and wherein the directive further includes a directive to determine a type for the classification;
trigger the LLM to operate on the LLM prompt, wherein, as a result of said operating, the LLM generates output comprising a specific classification for the candidate code snippet and a specific type for the specific classification;
format the output of the LLM into a data extension file that is consumable by the static analysis tool; and
include the data extension file in a corpus of data extension files that are consumable by the static analysis tool.
12 . The computer system of claim 11 , wherein the prompt is generated in real-time.
13 . The computer system of claim 11 , wherein the context includes information obtained from a source that is external to the code repository.
14 . The computer system of claim 11 , wherein the prompt is a fillable, pre-generated template.
15 . The computer system of claim 11 , wherein a size of the prompt is limited by a threshold size.
16 . The computer system of claim 11 , wherein the context includes a plurality of few-shot examples, and wherein a set of the few-shot examples are designed to encourage the LLM to assign negative classification to certain types of candidate code snippets.
17 . The computer system of claim 16 , wherein one type of candidate code snippet is structured to encourage the LLM to classify a particular candidate code snippet as not being a source or a sink.
18 . A method for implementing an artificial intelligence (Al) threat modeling tool, said method being performed by a computing service, wherein the computing service is one of a cloud service operating in a cloud environment, a local service operating on a local device, or a hybrid service that includes a cloud component operating in the cloud environment and a local component operating on the local device, said method comprising:
using a static analysis tool to extract a candidate code snippet from a code repository, wherein the candidate code snippet is identified by the static analysis tool as potentially being a security relevant code element;
using the static analysis tool to generate additional context associated with the candidate code snippet;
generating a large language model (LLM) prompt, which is structured to include the candidate code snippet, the context, and a directive to assign a classification to the candidate code snippet, said classification including a source classification, a sink classification, a sanitizer classification, or a flow step classification;
triggering the LLM to operate on the LLM prompt, wherein, as a result of said operating, the LLM generates output comprising a specific classification for the candidate code snippet; and
formatting the output of the LLM into a data extension file that is consumable by the static analysis tool.
19 . The method of claim 18 , wherein the prompt is a batch prompt that includes multiple candidate code snippets.
20 . The method of claim 18 , wherein the prompt further includes natural language instructions.