LLM-driven threat modeling

Abstract

Techniques for implementing an AI threat modeling tool are disclosed. A static analysis tool is used to extract candidate code snippets from a code repository. The candidate code snippets are identified as potentially security-related code elements. The static analysis tool generates an additional context associated with the candidate code snippet. LLM cues are generated. The hint is configured to include a candidate code snippet, a context, and an indication to assign a classification to the candidate code snippet. The classification includes source classification, sink classification, purifier classification, or flow step classification. The LLM operates on the hint to generate an output that includes the particular classification of the candidate code snippets. The output is formatted as a data extension file that is consumable by a static analysis tool.

Description

BACKGROUND

[0001] GitHub CodeQL is a static code analysis tool that can be used to find vulnerabilities and other issues within a codebase. For example, CodeQL can be used to identify and present security alerts. These alerts can correspond to vulnerabilities that consist of a“source” (i.e., a portion of code that allows“tainted” input to enter a system), a“sink” (i.e., a location where tainted data is used in a potentially dangerous way), and zero or more“flow steps” connecting them (where each flow step is a single step in the process that allows the taint to flow from the source to the sink). These alerts can exclude non-vulnerable code for which all flow paths from the source to the sink pass through a“sanitizer” (i.e., a portion of code that ensures that any data passing through it is no longer tainted). More specifically, a sanitizer is an entity that removes taint from data. For example, if there is a flow path from the source to the sink, but there is a sanitizer in between, then that path does not constitute a vulnerability. As another example, a SQL injection can consist of a flow path of a username supplied by a user and a method that reads information about the user from a SQL database. A sanitizer can include a method that checks whether the username is one of a specified set of usernames that are allowed to access, and if not, the SQL query will not run.

[0002] One major pain point in analyzing code involves the process of mapping or“threat modeling” a framework library (e.g., an open source repository) to identify sources, sinks, sanitizers, and flow steps. Historically, this mapping process has been extremely labor intensive and often involves countless hours of human effort. Moreover, the code in open source repositories is also growing. Keeping up with this growth (e.g., in terms of performing manual mapping) is extremely challenging. Performing this mapping on private repositories is even more difficult because the work needs to be done by security experts within the organization that owns the private repository. The number of private repositories is also growing at a high rate.

[0003] The subject matter claimed herein is not limited to implementations that solve any disadvantages or that operate only in environments such as those described above. Rather, this background is provided only to illustrate one exemplary technology area where some embodiments described herein can be practiced. SUMMARY

[0004] In some aspects, the technology described herein relates to a method for implementing an artificial intelligence (AI) threat modeling tool, the method comprising: extracting, using a static analysis tool, 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; generating, using the static analysis tool, additional context associated with the candidate code snippet; generating a large language model (LLM) prompt structured to include the candidate code snippet, the context, and an indication of a classification to assign to the candidate code snippet, the classification comprising a source classification, a sink classification, a sanitizer classification, or a flow step classification, and wherein the indication optionally further includes an indication of a type of the classification; triggering the LLM to operate on the LLM prompt, wherein, as a result of the operation, the LLM generates an output comprising a specific classification of the candidate code snippet and optionally a specific type of the specific classification; optionally formatting the output of the LLM into a data extension file consumable by the static analysis tool; and optionally including the data extension file in a corpus of data extension files consumable by the static analysis tool.

[0005] In some aspects, the technology described herein relates to a computer system comprising: a processor system; and a storage system comprising instructions executable by the processor system to cause the computer system to: extract, using a static analysis tool, 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; generate, using the static analysis tool, additional context associated with the candidate code snippet; generate a large language model (LLM) prompt structured to include the candidate code snippet, the context, and an indication of a classification to assign to the candidate code snippet, the classification comprising a source classification, a sink classification, a sanitizer classification, or a flow step classification, and wherein the indication further includes an indication of a type of the classification; trigger the LLM to operate on the LLM prompt, wherein, as a result of the operation, the LLM generates an output comprising a specific classification of the candidate code snippet and a specific type of the specific classification; format the output of the LLM into a data extension file consumable by the static analysis tool; and include the data extension file in a corpus of data extension files consumable by the static analysis tool.

[0006] In some aspects, the technology described herein relates to a method for implementing an artificial intelligence (AI) threat modeling tool, the method comprising: extracting, using a static analysis tool, 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; generating, using the static analysis tool, additional context associated with the candidate code snippet; generating a large language model (LLM) prompt structured to include the candidate code snippet, the context, and an indication of a classification to assign to the candidate code snippet, the classification comprising 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 the operation, the LLM generates an output comprising a specific classification of the candidate code snippet; and formatting the output of the LLM into a data extension file consumable by the static analysis tool.

[0007] This summary is provided to introduce some concepts of what follows in a simplified form. This summary is not intended to identify key or essential features of the claimed subject matter, nor is it intended to be used to determine the scope of the claimed subject matter.

[0008] Additional features and advantages will be set forth in the description that follows, and in part will be apparent from the description, or can be learned by practice of the teachings herein. The features and advantages of the present application will be realized and attained by means of the instruments and combinations particularly pointed out in the appended claims. Features of the present application will become more fully apparent from the following description and appended claims, or can be learned by the practice of the application as described hereinafter. BRIEF DESCRIPTION OF DRAWINGS

[0009] In order to describe the manner in which the above-recited and other advantages and features can be obtained, a more particular description will be rendered by reference to specific embodiments thereof, which are illustrated in the appended drawings. Understanding that these drawings depict only typical embodiments and are not therefore to be considered to be limiting of its scope, the embodiments will be described and explained with additional specificity and detail through the use of the accompanying drawings in which:

[0010] Figure 1 An example architecture for implementing a threat modeling tool is illustrated.

[0011] Figure 2 An example process flow is illustrated that outlines how a threat modeling tool can be implemented.

[0012] Figure 3 A flow diagram of an example method for classifying potentially security relevant code elements is illustrated.

[0013] Figure 4FIG. illustrates an example computer system that can be configured to perform any of the disclosed operations. DETAILED DESCRIPTION

[0014] The disclosed embodiments automate the detection (and facilitate the fixing) of vulnerabilities or potentially security-relevant code elements (also referred to as “threat models”) within a codebase. Once security-relevant code elements are identified, a static analysis tool can combine this information with control flow analysis and data flow analysis to identify security vulnerabilities involving these code elements.

[0015] One purpose of the disclosed embodiments is to attempt to find all or a subset of security-relevant code elements within a given codebase. Historically, manual mappings have not been scaled to cover all relevant code elements in all commonly used frameworks and libraries, resulting in an unacceptably high false negative rate. On the other hand, the disclosed embodiments are configured to better scale to detect security-relevant code elements and are configured to better detect vulnerabilities by effectively casting a larger detection net. Doing so can result in an initial set of false positives that is increased (compared to traditional techniques), but this result is acceptable because the embodiments are able to perform additional filtering, heuristics, or manual analysis. The embodiments also facilitate improvements in feedback learning by AI applications by expanding the set of examples that can optionally be included in LLM prompts.

[0016] The embodiments are highly configurable and, depending on the desired recall or precision level, the embodiments can fully automate the detection process or at least semi-automate the detection process. Specifically, the embodiments are designed to detect sources, sinks, sanitizers, and potentially flow steps within a codebase.

[0017] The embodiments achieve the above purposes by implementing an artificial intelligence (AI)-driven threat modeling system that can beneficially extend the coverage of CodeQL or any static analysis tool. Specifically, the disclosed embodiments help CodeQL identify more security-relevant elements in a codebase, thereby improving CodeQL’s ability to detect security vulnerabilities in that codebase and other codebases that depend on it, and reducing the false negative rate of CodeQL queries that detect security vulnerabilities. The embodiments use a combination of CodeQL and a large language model (LLM) to automatically detect sources, sinks, sanitizers, and flow steps. The embodiments then optionally make this information available to CodeQL queries. The CodeQL queries then combine this information with existing code analysis capabilities and previously known security-relevant elements to identify security vulnerabilities in the code.

[0018] Threat modeling can be formulated as a classification problem. For example, given a code snippet (e.g., argument of a method), the classification problem involves how to classify this code snippet as a sink of a particular type, a source of a particular type, a flow step of a particular type, a sanitizer of a particular type, or not classify as any of the above.

[0019] Embodiments are advantageously designed to solve this classification problem. To this end (as mentioned above), embodiments first trigger the use of CodeQL queries that are designed to extract candidate code snippets from a given code repository and heuristically filter a subset of these candidates as irrelevant. For example, if it is desired to model all sinks in a given framework, embodiments extract all arguments of method declarations, filter out some arguments with a set of heuristics. On the other hand, if it is desired to model all sinks used by a given application, embodiments extract all arguments of external application programming interface (API) calls, again filtering out as many irrelevant items as possible. Coverage can be increased by including other structures (including fields, qualifiers, and return values).

[0020] Embodiments also use CodeQL to extract positive and negative examples for few-shot prompting of a large language model (LLM). To illustrate, positive examples can include sinks, sources, sanitizers, or pollution steps that have been manually modeled. Negative examples can include entities that are known never to be sinks, sources, sanitizers, or pollution steps, or entities that do not need to be modeled (e.g., such as exceptions or code snippets located in test files).

[0021] Embodiments then use the information extracted from CodeQL to formulate prompts for the LLM. Depending on the LLM used, embodiments construct appropriate prompts to convey the relevant code snippet and selected additional context associated therewith, as well as a description of the task and optionally a selection of few-shot examples. Embodiments trigger or prompt the model (e.g., LLM) to return a string that indicates how the LLM classifies the code element as a sink of a particular type, a source of a particular type, a flow step of a particular type, a sanitizer of a particular type, or not classify as any of the above.

[0022] Advantageously, the disclosed technology can be generalized to any type of LLM, including completion-based and chat-based models as well as domain-specific code completion models. This generalization can be performed by adapting the construction of the prompts. For example, other formats of prompts can be used, such as markdown tables, chat style, mock code, or any other prompt style. Similarly, any type of static analysis tool can be used.

[0023] In this regard, the disclosed embodiments bring many benefits, advantages, and practical applications to the field of code management, particularly the field of detecting security-related code elements and static application security testing. Specifically, embodiments can beneficially detect security-related code elements within a codebase in an automated manner. By performing the disclosed operations, embodiments not only improve the way code is developed, but also improve the way code is managed. These embodiments also significantly improve the security of the underlying computer system. That is, by detecting and facilitating the repair of security-related code elements, the operation of the computer system is secured, thereby enabling the computer system to operate in a more reliable, secure, and robust manner.

[0024] As noted above, the disclosed techniques can be used to improve security. However, these principles can also be used to improve overall code quality. For example, these principles can be used to ensure that developers follow established coding practices, or develop code in a manner that complies with management policies. In this sense, the disclosed techniques can be used not only in the context of security, but can be generalized to improve overall code quality.

[0025] Embodiments can identify patterns in code, and map or analyze these patterns to determine a quality level of the code. That is, embodiments can be used to map particular code within a library or framework. Embodiments can then generate various queries to determine how a candidate code snippet connects with other identified code snippets or patterns. Accordingly, these and many other benefits will now be described in greater detail in the remainder of this disclosure. Example architecture

[0026] Having just described some high-level benefits and advantages of the disclosed embodiments, attention will now turn to Figure 1 . Figure 1 A computing architecture 100 that can be used to implement the benefits described above is illustrated. Architecture 100 is shown to include a service 105.

[0027] As used herein, the term "service" refers to an automated program tasked with performing different actions based on input. In some cases, service 105 can be a deterministic service that operates completely given a set of inputs and without randomizing factors. In other cases, service 105 can be or include a machine learning (ML) or artificial intelligence engine. The ML engine enables service 105 to operate even in the face of randomizing factors.

[0028] As used herein, a reference to any type of machine learning or artificial intelligence can include any type of machine learning algorithm or device, convolutional neural network(s), multi-layer neural network(s), recurrent neural network(s), deep neural network(s), decision tree model(s) (e.g., decision trees, random forests, and gradient boosting trees), linear regression model(s), logistic regression model(s), support vector machine(s) (“SVM”), artificial intelligence device(s), or any other type of intelligent computing system. The machine learning algorithm(s) can be trained using any number of training data (and possibly later refined) to dynamically perform the disclosed operations.

[0029] The service 105 is shown to include or be associated with a large language model (LLM) 110. The LLM 110 can represent the machine learning engine or artificial intelligence described above. The LLM 110 is a type of neural network that uses layers of nodes in a probabilistic manner. In response to a prompt, the LLM 110 generates probabilities for words to form various groupings of words. The LLM 110 can be a first-party LLM or a third-party LLM.

[0030] In some implementations, the service 105 is a cloud service that operates in a cloud environment 115. In some implementations, the service 105 is a local service that operates on a local device. In some implementations, the service 105 is a hybrid service that includes a cloud component that operates in the cloud and a local component that operates on a local device. The two components can communicate with each other.

[0031] The service 105 is also shown to include a static analysis tool 120 (e.g., CodeQL). The service 105 uses the static analysis tool 120 to access a codebase 125 and extract candidate 130 code snippets and specified information about the candidate from the codebase 125. Notably, the service 105 is designed to obtain an overextended or overly broad set of potential security-related elements (e.g., the candidate 130 is one of these elements).

[0032] In some cases, various filtering operations or pre-filtering or other heuristic operations can be performed in an attempt to identify security-related code elements. These operations are designed to eliminate a subset of false positives that are captured during an initial operation of analyzing the codebase 125. In some scenarios, pre-filtering can be performed to reduce the number of candidates, thereby reducing LLM computational costs. Another example of pre-filtering involves reducing the number of false positives, some of which can be more easily identified by heuristic filters during static analysis.

[0033] As one example, code snippets that have already been mapped need not be remapped by the service 105. Thus, these code snippets can optionally be filtered from the analysis. As another example, there can be cases where certain frameworks are known not to include vulnerabilities. Calls to APIs in these frameworks can also be filtered. As another example, there can be cases where certain test code is irrelevant and can be filtered. Thus, the service 105 can employ a set of filters to reduce the amount of data, and thus also the number of subsequent calls to the LLM 110, whose task is to classify these potentially security-relevant elements. Thus, by performing these filtering and heuristic operations, embodiments are able to intelligently reduce the amount of data that will be fed to the LLM 110, thus also improving the operational efficiency and computational cost of the computing system.

[0034] The service 105 (e.g., via the static analysis tool 120) also accesses various examples (and the relevant information or features of each example), or more generally, accesses the context 135 associated with the candidate 130. The context 135 can include any type of data. As one example, the context 135 can be the content pulled from a selected number of lines of code before the line or lines that include the candidate 130. As another example, the context 135 can be the content pulled from a selected number of lines of code after the line or lines that include the candidate 130. Thus, the context 135 should be interpreted broadly and can include any type of supplemental data, including few-shot examples. For example, the context 135 can include positive and negative examples or other features. As another example, if the candidate is a method parameter, the context can include the parameter name as well as the name, signature, type or class name, and package of the method.

[0035] The context 135 can include the docstring of the method (or library defining the method) that includes the candidate 130. Method parameter information can also be included in the context 135. Code comments, code documentation, and other information about libraries, APIs, or methods related to the code base 125 can also be included in the context 135.

[0036] It should be mentioned that in some implementations, how the static analysis tool 120 is able to extract the context 135. For example, in scenarios where the static analysis tool 120 employs CodeQL form, CodeQL is able to extract or otherwise obtain the context 135 of the candidate 130 from the code base 125. It should also be mentioned that the context can generally be extracted from the code base itself and / or other sources, such as the internet. In some scenarios, the code base 125 and possibly other sources of context can operate as inputs to the service 105. The extraction of the candidate 130 and the context 135 can optionally be included as part of the service 105.

[0037] In some scenarios, the service 105 (particularly the static analysis tool 120) has access to the underlying code, while in other scenarios, the tool can not have access to the underlying code. Instead, the tool can have access to the code that calls the method that includes the candidate code snippet, but not the code of the method itself. Thus, in some scenarios, embodiments extract candidate code snippets from calls to external API calls in an application or set of programs, while in other scenarios, embodiments extract code snippets from public methods in the framework / library itself, framing the scenario by framework type (e.g., embodiments can query open source frameworks for API calls defined within the framework and data expected to be obtained for the context 135).

[0038] In this way, embodiments are able to identify a rich set of semantic information about the candidate 130. In other words, information that is syntactically and / or semantically related to the candidate 130 can be included in the context 135. Thus, as used herein, the term “context” should be interpreted in a broad manner and can include any data, metadata, or any other information pulled from the code base 125 and / or pulled from sources external to the code base 120 (but including data related to the candidate 130).

[0039] Information about how the method is called can also be included in the context 135. Additionally, external data can be included in the context 135. For example, release notes or change notes can be queried to determine if they include valuable contextual information to be included in the context 135.

[0040] The service 105 then formulates a prompt 140. The prompt 140 is constructed by the service 105 to include the context 135 and the current candidate 130 to be classified. In some implementations, the prompt 140 is a fillable template, where data can be filled into fields of the template. In other implementations, the prompt 140 is dynamically and in real-time generated by the service 105. In this way, the service 105 reformulates the classification problem as a text generation and question answering problem, such that the text completion or chat response provided by the LLM 110 corresponds to the classification of the current candidate 130.

[0041] Some embodiments can limit the size of the prompt 140 based on technical limitations, design considerations, or operational considerations (e.g., cost) relative to the LLM 110.

[0042] The prompt 140 can be further constructed to include few-shot examples. These few-shot examples operate to guide the LLM 110 to generate the desired type of output.

[0043] As one example, few-shot examples can be generated and included in the prompt 140, where the few-shot examples include definitions or examples of what constitutes a source or what constitutes a sink or what constitutes a flow step or what constitutes a sanitizer. Negative examples can also be provided, such as examples that do not constitute a source, sink, sanitizer, or flow step.

[0044] Examples can also be provided for different types of sinks (e.g., this is a sink of the SQL type), similar to sources, sanitizers, and flow steps. Few-shot examples are particularly beneficial for fine-grained distinctions that will be used to help guide the analysis trajectory of the LLM 110. Few-shot examples are particularly beneficial for addressing false positives and false negatives that can be ambiguous or difficult to clearly distinguish, or that have not yet reached community consensus.

[0045] Few-shot examples can also be used to filter data. For example, even though a particular code unit can technically be a sink, a developer can not be interested in that code unit because the code is test code. If that is the case, few-shot examples can be prepared so that the LLM 110 will be prevented from negatively classifying that data unit. Another potential sink of disinterest can be a sink that is deeply buried within a hierarchy, and thus not exposed outwardly (and not easily used maliciously).

[0046] The static analysis tool 120 can be configured to identify these few-shot examples. For example, CodeQL can be configured to query any number of databases that contain previously modeled code snippets. CodeQL can then extract relevant examples (positive examples and / or negative examples) from these example databases, which can operate as few-shot examples in the prompt 140.

[0047] Optionally, the static analysis tool 120 can use previously recorded data to generate few-shot examples, and / or the static analysis tool 120 can use real-time data (e.g., code currently included in the codebase 125) to generate few-shot examples. Few-shot examples can also optionally be manually selected. Optionally, some few-shot examples can be hardcoded in the prompt. In the prompt, a rationale or explanation can be provided to inform the LLM 110 why an example is a positive example or a negative example.

[0048] Few-shot examples can also be used to guide the format of the LLM output. Ultimately, the task of the service 105 will be to parse the output of the LLM so that it can be integrated back into the static analysis tool 120 (e.g., in the form of a data expansion file 150). Few-shot examples can be used to guide the formatting of the LLM output. The data expansion file 150 can be considered a corpus of reference data used by the static analysis tool 120 to enable it to perform its analysis operations. The newly generated data expansion file 150 can be included in an existing corpus 155 of data expansion files, some of which can include manually modeled information and some of which can have been generated via automated techniques such as described herein.

[0049] In some scenarios, a single prompt is generated for each candidate code snippet. In other scenarios, a single prompt can batch multiple candidate code snippets and trigger the LLM to generate a classification for each candidate in the batch. In some cases, multiple prompts can be generated for a single candidate, such as in a feedback prompt scenario.

[0050] For example, a first prompt can be generated for a particular candidate, and the LLM can attempt to generate a classification. If that classification is incorrect, flawed, or needs further refinement, a feedback / second (or third, fourth, etc.) prompt can be generated with additional data and reasoning about why additional processing is requested.

[0051] The prompt 140 can also include some natural language instructions that describe the purpose or goal of the LLM 110. For example, the natural language instructions can describe the desired outcome of the LLM 110 and a statement about how to format the output, such as a statement that the LLM is to provide types for a sink, a source, a sanitizer, and / or a flow step in a specified format and the rationale for that determination.

[0052] The LLM 110 then operates based on the prompt 140. The task of the LLM 110 is to classify the candidate 130 as a sink, a source, a flow step, a sanitizer, or otherwise. Additionally, the task of the LLM 110 is to determine the type of sink or the type of source or the type of flow step or the type of sanitizer.

[0053] As an example, a “type” of sink can be a cross-site scripting sink. This type of sink is a sink that writes data to a web page. Another type of sink can include a SQL injection sink or a query injection sink, both of which generally involve reading, writing, or deleting data in a database. Another example of a sink type is a command injection sink, where user-provided content will affect a subsequently executed code snippet. Another example type is a path injection sink, which is a file system path that will be read, written, opened, or deleted. Of course, other types of sinks exist, and the above enumerated should be considered not exhaustive.

[0054] Examples of the type of source can be the origin of remote untrusted data that an attacker can attempt to send into the application. As another example, the type of source can be user input into the application. The source type can be local or remote. The few-sample examples mentioned previously can include example types of sources, sinks, sanitizers, and flow steps. Thus, the hints can include examples of types, and the hints can specify a goal for the LLM to provide types.

[0055] In some implementations, the enumeration of few-sample examples can be structured in a way that limits the LLM 110 such that the LLM 110 is only allowed to select a classification from the examples provided in the hints. In other implementations, the enumeration of few-sample examples is structured to provide additional flexibility to the LLM 110 to potentially select a classification that is not within the range of few-sample examples.

[0056] The service 105 then uses the resulting classification 145, which can also include a type 145A for the classification, to generate data extension files 150. The service 105 can then optionally feed these files 150 back into the static analysis tool 120 (e.g., into the corpus 155) to augment the default analysis capabilities of the tool. The type 145A and the classification 145 can be included as part of the data extension files 150. In some scenarios, these files can be automatically added to the static analysis tool, or they can first be manually verified and categorized.

[0057] That is, the static analysis tool 120 (e.g., CodeQL) is designed to consume data extension files that are kept in the corpus 155, and these files can optionally be written in YAML (another markup language) or possibly another language. These data extension files can provide additional information about API functions to the static analysis tool 120 in a specific format that the tool already understands. The static analysis tool 120 can be equipped with a large collection of these data extension files (e.g., as shown in the corpus 155), and embodiments are able to use the AI-generated classifications (e.g., the classification 145) to further augment the modeling capabilities of the tool.

[0058] In practice, the service 105 provides a pathway for new classifications (e.g., classification 145) to be beneficially incorporated into the data extension files (included in corpus 155) of the static analysis tool 120. Thus, the static analysis tool 120 can use both manual and AI-generated modeling to identify more security-related code and security vulnerabilities, which in turn can be surfaced to developers. In parallel, embodiments optionally enable customers to use the same automated modeling techniques to create custom data extensions that are applied to their own private repositories or public repositories they rely on. This can be implemented through a service 105 providing a user experience (UX) (optionally in an integrated development environment (IDE) or in a browser) such that users can run the service 105 (such as by implementing an automated model operation), categorize the resulting classifications, and include those classifications in the static analysis tool 120.

[0059] Figure 2 An example process flow 200 is shown that can be performed by the service 105 of Figure 1 Generally, the process flow 200 includes two steps or two stages, as shown in stage one 205 and stage two 210.

[0060] Stage one 205 involves the use of a static analysis tool 215, which represents the static analysis tool 120 of Figure 1 During stage one 205, the static analysis tool 215 is used to analyze a codebase, and from within that codebase, one or more candidate code snippets are identified that can potentially be security-related code elements.

[0061] Stage two 210 generally involves pre-processing these candidates, and optionally filtering or reducing redundancy from them, generating appropriate hints for them (including context associated with the snippet and optionally a few sample examples), and feeding those hints to an LLM 220. The LLM 220 represents the LLM 110 from Figure 1 The LLM 220 operates to classify these code snippets as particular types of vulnerabilities or particular types of security-related code elements. These classifications can then be analyzed and post-processed (optionally categorized manually), and fed back to the static analysis tool 215, whereby the tool can improve its subsequent identification processes.

[0062] In some cases, post-processing can be performed. As one example, assume that multiple LLM calls are made with respect to the same candidate code snippet. However, there can be cases where the context of the code snippet differs in each or at least some of the multiple calls. Some embodiments are configured to post-process the results generated by the LLMs to attempt to resolve potential inconsistencies or differences that can exist between the different outputs produced by the LLMs. For example, a voting mechanism can be implemented, voting among the various different outputs, to attempt to reach a single solution. In some implementations, different weightings or expert knowledge can be incorporated, such that the selected result is not strictly based on a simple majority choice.

[0063] Post-processing can also include logic to format the output in a manner that is understandable to the static analysis tool. Recall that the classification output generated by the LLMs is typically fed back to the static analysis tool in the form of one or more data extension files (which will be included in the corpus of reference data used by the static analysis tool). Thus, post-processing can be performed to ensure that the format of the data is sufficient for the static analysis tool to operate on the output. Thus, in various embodiments, post-processing refinement can be performed on the output of the LLMs.

[0064] Optionally, the task of the LLM can be to format its output in a desired manner. As another option, separate post-processing can be performed to reformat the output of the LLM in a desired manner. In some implementations, a combination of these two options can also be performed.

[0065] As mentioned previously, embodiments can also require the LLM to provide a rationale for its classification decision. In effect, embodiments can employ chain-of-thought reasoning with respect to the LLM. Requiring the LLM to provide a rationale has been shown to significantly improve the classification output generated by the LLM in experiments with the disclosed embodiments. Thus, by requiring the LLM to perform chain-of-thought reasoning, significant advantages can be realized.

[0066] The resulting extended data file includes various threat models that static analysis tools can use to perform their operations. However, it's important to note how the models in the extended data file differ from the output of the queries initially run by the static analysis tool and constituting the candidate code snippets (i.e., those potentially security-related code elements). Queries designed to find potentially security-related code elements are overly broad to capture all or at least the vast majority of potentially security-related code elements present in the codebase. By design, they output many candidates that are actually not security-related. On the other hand, the models generated by the service and included in the extended data file are designed to be a significantly smaller subset of code elements, filtering out most non-security-related elements. Furthermore, these models often include types not present in the candidates generated by the static analysis queries. The ultimately detected threat models can be used to improve the analytical capabilities of the static analysis tool.

[0067] It should also be noted that queries for detecting candidate code snippets differ from queries written for detecting vulnerabilities in source code. The former is designed to detect potential sources, sinks, purifiers, and flow steps; the latter is designed to connect such units together to detect or identify vulnerabilities (e.g., by connecting identified sources to identified sinks via a series of contamination steps). The former is part of the disclosed embodiments, while the latter can use the output of the embodiments just like manually modeled sources, sinks, purifiers, and contamination steps, and is not specific to the disclosed embodiments. Therefore, it should be understood how the disclosed static analysis tool uses two different sets of queries. Example method

[0068] The following discussion now refers to a variety of methods and method actions that can be performed. Although method actions may be discussed in a specific order or illustrated in a flowchart as occurring in a specific order, a specific order is not required unless specifically stated otherwise, or because an action depends on another action that is performed before that action is performed.

[0069] Now let's turn our attention to... Figure 3 , Figure 3 The diagram illustrates a flowchart of an example method 300 for implementing an artificial intelligence (AI) threat modeling tool. Method 300 can be... Figure 1 The architecture is implemented within 100. Furthermore, method 300 can be implemented by... Figure 1 Service 105 is implemented.

[0070] Method 300 includes using static analysis tools (e.g., from...) Figure 1The action 305 includes an action of extracting a candidate code snippet (candidate 130) from a code repository (e.g., code base 125) using a static analysis tool 120). The candidate code snippet is identified by the static analysis tool as potentially being a security-relevant code element.

[0071] In some scenarios, the candidate code snippet is one of a plurality of candidate code snippets identified and extracted by the static analysis tool. Optionally, a pre-filtering operation is performed on these candidate code snippets using a defined set of parameters to reduce the number of candidate code snippets.

[0072] The action 310 includes an action of generating additional context (e.g., context 135) associated with the candidate code snippet using the static analysis tool. In some cases, the context includes one or more few-shot examples. These few-shot examples can include at least one of: positive examples of corresponding sinks, positive examples of corresponding sources, positive examples of corresponding sanitizers, or positive examples of corresponding flow steps. These few-shot examples can also include at least one of: negative examples of corresponding sinks, negative examples of corresponding sources, negative examples of corresponding sanitizers, or negative examples of corresponding flow steps. In some scenarios, some of the context can be extracted using other tools, such as web searching for documentation of the candidate API call. Positive and negative examples of sanitizers can also be provided.

[0073] The context can also include documentation and / or code comments associated with the candidate code snippet. The context can also include content extracted from a selected number of lines of code before or after the line of code that includes the candidate code snippet. In some cases, the context also includes information about how the function that includes the candidate code snippet is called.

[0074] Optionally, the context can include information obtained from sources outside the code repository, such as pulled from forums, social media sites, release notes, etc. As another option, the context can include a plurality of few-shot examples. The set of few-shot examples can be used to influence the classification assigned to certain types of candidate code snippets by the LLM. For example, few-shot examples can be constructed to operate as filters or pseudo-filters to prevent certain types of candidate code snippets from being classified as sinks, sources, sanitizers, or flow steps, thereby causing the LLM to label these snippets as none of the above. As one example, negative few-shot examples can encourage the LLM to assign a negative classification to non-user-facing candidate code snippets. Further, in some scenarios, a type of candidate code snippet is constructed to encourage the LLM to classify a particular non-user-facing candidate code snippet as either non-source or non-sink.

[0075] Action 315 includes generating a large language model (LLM) prompt (e.g., prompt 140). The prompt is structured to include the candidate code snippet, the context, and an implicit or explicit indication of the classification to assign to the candidate code snippet. The classification includes a source classification, a sink classification, a sanitizer classification, or a flow step classification. The indication can also include an indication of the type of classification determined. Optionally, the prompt is generated in real-time. As another option, the prompt can be a fillable, pre-generated template. In some instances, the size of the prompt is limited by (or subject to) a threshold size.

[0076] In some cases, the prompt is a batch prompt that includes multiple candidate code snippets. In some cases, the prompt includes a single candidate code snippet. Optionally, the prompt can include additional natural language instructions to the LLM.

[0077] Action 320 includes triggering the LLM to operate on the LLM prompt. As a result of its operation, the LLM generates an output that includes a specific classification of the candidate code snippet (e.g., classification 145) and a specific type of the specific classification (e.g., type 145A).

[0078] Action 325 includes formatting the output of the LLM into a data augmentation file (e.g., data augmentation file 150). The file is consumable by the static analysis tool. The process of formatting the output of the LLM into a data augmentation file can include parsing the output, and organizing the parsed output into a format that is consumable by the static analysis tool.

[0079] Action 330 subsequently includes inserting or including the data augmentation file into a corpus of data augmentation files. The corpus is consumable by the static analysis tool, and is operated on to improve or enhance the analysis capabilities of the static analysis tool. In some cases, inserting the data augmentation file into the corpus of data augmentation files can be performed manually and / or after a triage event.

[0080] Accordingly, the disclosed technology is beneficially configured to avoid having to directly ask the LLM to classify arbitrary code snippets as vulnerable or non-vulnerable, which can result in low precision results. Instead, embodiments use the LLM to classify the security relevance of each code snippet with a richer set of possible classifications (e.g., a source of untrusted data, a sink of a particular vulnerability type, a flow step that propagates data, or a sanitizer that removes pollution), and then pass those classifications back to a highly precise static analysis tool that can consume those classifications to extend its own analysis coverage capabilities, rather than directly passing the classifications to a software developer.

[0081] In this regard, the static analysis tool is used in two ways. First, it is used to identify candidate code elements for classification by the LLM along with contextual information. Second, it is used to consume the classifications generated by the LLM to augment its own static analysis capabilities.

[0082] Any type of static analysis tool can be used, and CodeQL is just one example. In fact, any static analysis tool that can reason about the syntax and semantics of source code, precisely identify candidate code elements, and consume classifications can be used. The static analysis tool is beneficially used to extract candidates and few-shot examples, thereby beneficially limiting the number of calls to the LLM. This technique is in contrast to providing all code to the LLM. Embodiments also provide a principled framework for eliminating false positives via candidate filtering and presentation of various negative examples. Using a pre-trained LLM also advantageously eliminates the need to build large training sets and build or maintain custom models.

[0083] Building classifications as code completion problems allows embodiments to use models that are fine-tuned on code that perform as well as larger general-purpose models. However, the disclosed approach generalizes to completion and chat-based general-purpose models by adapting the prompting strategy. Providing explanations in the form of annotations to few-shot examples improves both accuracy via the low-cost think-link approach and transparency by providing explanations for the model’s classification decisions. Example computer / computer system

[0084] Attention is now directed towards Figure 4 , Figure 4 FIGURE 1 illustrates an example computer system 400 that can include and / or be used to perform any of the operations described herein. For example, the computer system 400 can implement the method 300 of Figure 3 Similarly, the computer system 400 can implement the architecture 100 of Figure 1

[0085] The computer system 400 can take various different forms. For example, the computer system 400 can be implemented as a tablet computer, a desktop computer, a laptop computer, a mobile device, or a standalone device such as those described throughout this disclosure. The computer system 400 can also be a distributed system that includes one or more connected computing components / devices in communication with the computer system 400.

[0086] In its most basic configuration, the computer system 400 includes various different components. Figure 4 As shown, the computer system 400 includes a processor system 405 that includes one or more processors (also referred to as “hardware processing units”) and a storage system 410.

[0087] ​With respect to the processor(s) of the processor system 405, it is understood that the functionality described herein can be performed, at least in part, by one or more hardware logic components. For example, and without limitation, illustrative types of hardware logic components that can be used include Field-Programmable Gate Arrays ("FPGAs"), Program-Specific or Application-Specific Integrated Circuits ("ASICs"), Program- Specific Standard Products ("ASSPs"), System-on-a-Chip ("SOCs"), Complex Programmable Logic Devices ("CPLDs"), Central Processing Units ("CPUs"), Graphics Processing Units ("GPUs"), or any other type of programmable hardware.

[0088] As used herein, the terms "executable module", "executable component", "component", "module", "tool", or "engine" can refer to a hardware processing unit or to a software object, routine, or method that can be executed on the computer system 400. The different components, modules, engines, and services described herein can be implemented as objects or processors executing on the computer system 400, for example, as separate threads.

[0089] The storage system 410 can include physical system memory, which can be volatile, nonvolatile, or some combination of the two. The term "memory" can also be used herein to refer to non-volatile mass storage, such as physical storage media. If the computer system 400 is distributed, the processing, storage, and / or memory capabilities can be distributed as well.

[0090] The storage system 410 is shown to include executable instructions 415. The executable instructions 415 represent instructions that are executable by the processor(s) of the computer system 400 to perform the disclosed operations, such as those described in the various methods.

[0091] The disclosed embodiments can include or utilize special-purpose or general-purpose computers including computer hardware, such as one or more processors and system memory, as discussed in greater detail below. Embodiments also include physical and other computer-readable media for carrying or having computer-executable instructions and / or data structures stored thereon. Such computer-readable media can be any available media that can be accessed by a general purpose or special purpose computer system. Computer-readable media that store computer-executable instructions in data form are "physical computer storage media" or "hardware storage devices." Moreover, computer-readable storage media, including physical computer storage media and hardware storage devices, exclude signals, carriers, and transitory signals. On the other hand, computer-readable media that carry computer-executable instructions are "transmission media", and include signals, carriers, and transitory signals. Thus, the present embodiments can include at least two distinct kinds of computer-readable media: computer storage media and transmission media.

[0092] Computer storage media (also referred to as "hardware storage devices") are computer- readable hardware storage devices, such as RAM, ROM, EEPROM, CD-ROM, solid state drives ("SSDs") that are based on RAM, Flash memory, phase-change memory ("PCM"), or other types of memory, or other optical disk storage, magnetic disk storage or other magnetic storage devices, or any other hardware storage devices that can be used to store desired program code means in the form of computer-executable instructions, data, or data structures and that can be accessed by a general purpose or special purpose computer.

[0093] The computer system 400 can also be connected to external sensors (e.g., one or more remote cameras) or devices via the network 420 (via wired or wireless connection). For example, the computer system 400 can communicate with any number of devices or cloud services to obtain or process data. In some cases, the network 420 itself can be a cloud network. Moreover, the computer system 400 can also be connected to remote / separate computer system(s) over one or more wired or wireless networks, which are configured to perform any of the processes described with respect to the computer system 400.

[0094] A "network" (such as the network 420) is defined as one or more data links and / or data switches that enable the transport of electronic data between computer systems, modules, and / or other electronic devices. When information is transferred or provided over a network (either hardwired, wireless, or a combination of hardwired and wireless) to a computer, the computer properly views the connection as a transmission medium. The computer system 400 will include one or more communication channels that are used to communicate with the network 420. Transmission media include a network that can carry data or desired program code means in the form of computer-executable instructions or in the form of data packets in a modulated data signal. The computer-executable instructions can be accessed by a general purpose or special purpose computer. Combinations of the above should also be included within the scope of computer-readable media.

[0095] At the various computer system components, program code means in the form of computer-executable instructions or data structures can be transferred automatically from transmission media to computer storage media (or vice versa). For example, computer-executable instructions or data structures received by way of network or data link can be buffered in RAM within the network interface module (e.g., a network interface card or "NIC"), then eventually transferred to computer system RAM and / or to less volatile computer storage media at the computer system. Accordingly, it should be understood that computer storage media can be included in computer system components that also (or even primarily) use transmission media.

[0096] Computer-executable (or computer-interpretable) instructions comprise, for example, instructions which cause a general purpose computer, special purpose computer, or special purpose processing device to perform a certain function or group of functions. The computer executable instructions can be, for example, binaries, intermediate format instructions such as assembly language, or even source code. Although the subject matter has been described in language specific to structural features and / or methodological acts, it is to be understood that the subject matter defined in the appended claims is not necessarily limited to the described features or acts described above. Rather, the described features and acts are disclosed as example forms of implementing the claims.

[0097] Those skilled in the art will appreciate that embodiments can be practiced in network computing environments with many types of computer system configurations, including personal computers, desktop computers, laptop computers, message processors, hand-held devices, multi-processor systems, microprocessor-based or programmable consumer electronics, network PCs, minicomputers, mainframe computers, mobile telephones, PDAs, pagers, routers, switches, and the like. Embodiments can also be practiced in distributed system environments where local and remote computer systems, which are linked (either by hardwired data links, wireless data links, or by a combination of hardwired and wireless data links) through a network, each perform tasks (e.g., cloud computing, cloud services, etc.). In a distributed system environment, program modules can be located in both local and remote memory storage devices.

[0098] This application can be implemented in other specific forms without departing from its spirit or central characteristics. The described embodiments are to be considered in all respects only as illustrative and not restrictive. The scope of the application is, therefore, indicated by the appended claims, rather than by the foregoing description. All changes that come within the meaning of equivalency of the claims are intended to be embraced within the scope of the claims.

Claims

1. A method (300) for implementing an artificial intelligence (AI) threat modeling tool, the method comprising: extracting (305), using a static analysis tool (120), 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; generating (310), using the static analysis tool, additional context (135) associated with the candidate code snippet; generating (315) a large language model (LLM) prompt (140) structured to include the candidate code snippet, the context, and an indication of an assignment of a classification to the candidate code snippet, the classification comprising a source classification, a sink classification, a sanitizer classification, or a flow step classification, and wherein the indication further includes an indication of a determination of a type for the classification; triggering (320) the LLM to operate on the LLM prompt, wherein, as a result of the operation, the LLM generates an output comprising a specific classification for the candidate code snippet and a specific type for the specific classification; formatting (325) the output of the LLM into one or more data extension files consumable by the static analysis tool; and including (330) the data extension files in a corpus of data extension files consumable by the static analysis tool.

2. The method of claim 1, wherein the context comprises one or more few-shot examples.

3. The method of claim 2, wherein the one or more few-shot examples comprise 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 comprise 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 comprises a document associated with the candidate code snippet.

6. The method of claim 1, wherein the context further comprises content extracted from a selected number of lines of code preceding a line of code comprising the candidate code snippet.

7. The method of claim 1, wherein the context further comprises content extracted from a selected number of lines of code following a line of code comprising the candidate code snippet. parsing the output, and organizing the parsed output into a format consumable by the static analysis tool.

8. The method of claim 1, wherein formatting the output of the LLM into the data extension file comprises:

9. The method of claim 1, wherein the candidate code snippet is one of a plurality of candidate code snippets 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 about a particular classification that is provided to a few-shot example, the few-shot example also being included in the prompt.

11. A computer system (400) comprising: a processor system (405); and a storage system (410) comprising instructions executable by the processor system to cause the computer system to: extract (305), using a static analysis tool (120), candidate code snippets from a code repository, wherein the candidate code snippets are identified by the static analysis tool as potentially being security-relevant code elements; generate (310), using the static analysis tool, additional context (135) associated with the candidate code snippets; generate (315) a large language model (LLM) prompt (140) structured to include the candidate code snippets, the context, and an indication of a classification to assign to the candidate code snippets, the classification comprising a source classification, a sink classification, a sanitizer classification, or a flow step classification, and wherein the indication further includes an indication of a type determined for the classification; trigger (320) the LLM to operate on the LLM prompt, wherein, as a result of the operation, the LLM generates an output comprising a specific classification for the candidate code snippets and a specific type for the specific classification; format (325) the output of the LLM into a data extension file consumable by the static analysis tool; and include (330) the data extension file in a corpus of data extension files 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 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 the set of few-shot examples are designed to encourage the LLM to assign a 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 that is not user-facing as not being: a source or a sink.

18. A method (300) for implementing an artificial intelligence (AI) threat modeling tool, the method comprising: extracting (305), using a static analysis tool (120), candidate code snippets from a code repository, wherein the candidate code snippets are identified by the static analysis tool as potentially being security-relevant code elements; generating (310), using the static analysis tool, additional context (135) associated with the candidate code snippets; generating (315) a large language model (LLM) prompt (140) structured to include the candidate code snippet, the context, and an indication of a classification assigned to the candidate code snippet, the classification comprising a source classification, a sink classification, a sanitizer classification, or a flow step classification; triggering (320) the LLM to operate on the LLM prompt, wherein, as a result of the operation, the LLM generates an output comprising a specific classification for the candidate code snippet; and formatting (325) the output of the LLM into a data extension file consumable by the static analysis tool.

19. The method of claim 18, wherein the prompt is a batch prompt comprising a plurality of candidate code snippets.

20. The method of claim 18, wherein the prompt further comprises a natural language instruction.

Images