An automatic penetration testing method for large language models based on multi-agents

Abstract

The present invention belongs to the field of artificial intelligence and network security technology, and specifically relates to a large language model automatic penetration testing method based on multiple agents, including building a penetration testing solution knowledge base and a penetration testing tool library; establishing five large model agents, including a knowledge retrieval agent, a task planning agent, a task scheduling agent, a tool calling agent, and a task summarizing agent. The present invention uses the architecture of multi-agent collaboration and combines the advantages of RAG technology to effectively reduce the hallucination problem that may occur in the decision-making process of the large language model when processing complex penetration testing tasks, thereby improving the completion rate and accuracy of penetration testing tasks.

Description

Technical Field

[0001] The present invention belongs to the field of artificial intelligence and network security technology, and specifically relates to an automatic penetration testing method of a large language model based on multi-agents. Background Art

[0002] Penetration Testing is a method of evaluating the security of a computer system, network, or application by simulating an attack. The purpose of penetration testing is to discover potential security vulnerabilities and provide recommendations on how to strengthen the system to prevent potential attacks. It can help identify the system's attack vulnerabilities and estimate its vulnerability.

[0003] At present, penetration testing still mainly relies on professionals and has a low degree of automation. Although penetration testing technology is constantly developing and improving, there are still some shortcomings, some of which include: ① Over-reliance on manual experience: Penetration testing usually requires experienced professionals who understand various attack techniques and vulnerability exploitation methods. This makes the penetration test results likely to be limited by individual experience and skill levels. ② Limited rule writing: Automated penetration testing tools often rely on pre-defined rule sets that are used to detect potential vulnerabilities. However, these rules are usually written based on known vulnerabilities and attack methods, which makes it difficult to cover all possible situations. New attack techniques may not be effectively detected by these rules. ③ Lack of comprehensiveness: Penetration testing usually involves multiple attack surfaces, including networks, applications, physical devices, etc. However, current penetration testing methods and tools may still appear to be not comprehensive enough when integrating these aspects, resulting in limitations in testing. Summary of the invention

[0004] The purpose of the present invention is to provide a large language model automatic penetration testing method based on multi-agents to improve the efficiency, accuracy and success rate of automatic penetration testing.

[0005] To achieve the above object, the technical solution adopted by the present invention is:

[0006] A multi-agent-based large language model automatic penetration testing method, the multi-agent-based large language model automatic penetration testing method comprising:

[0007] Build a penetration testing solution knowledge base and a penetration testing tool library;

[0008] Establish five large-scale model agents, including knowledge retrieval agent, task planning agent, task scheduling agent, tool calling agent and task summary agent;

[0009] Receive an input task description, call the knowledge retrieval agent, and retrieve a knowledge corpus in the penetration testing solution knowledge base according to the task description;

[0010] Calling the task planning agent to convert the task description and the knowledge corpus into a first prompt word, and obtaining a penetration test step list according to the first prompt word;

[0011] For each penetration test step in the penetration test step list, call the task scheduling agent, convert the tool output results of the current penetration test step and all penetration test steps before the current penetration test step into the second prompt word, and obtain the operation type corresponding to the current penetration test step according to the second prompt word. If the operation type is tool call, call the tool call agent; if the operation type is execution summary, call the task summary agent;

[0012] The tool calling agent retrieves the corresponding program interface in the penetration testing tool library according to the current penetration testing step, combines the current penetration testing step, the previous penetration testing step, and the execution result of the previous penetration testing step into a third prompt word, obtains the tool output result corresponding to the current penetration testing step according to the third prompt word, and inputs the tool output result into the task scheduling agent;

[0013] The task summary agent converts the task description and the execution results of all penetration test steps into a fourth prompt word, and outputs the automatic penetration test result according to the fourth prompt word.

[0014] Several optional methods are also provided below, but they are not intended to be additional limitations on the above-mentioned overall solution, but are merely further supplements or preferences. Under the premise that there are no technical or logical contradictions, each optional method can be combined with the above-mentioned overall solution separately, and multiple optional methods can also be combined.

[0015] Preferably, the penetration testing solution knowledge base is constructed as follows:

[0016] Acquire penetration testing professional knowledge corpus through web crawling;

[0017] The knowledge corpus in the field of penetration testing is divided according to the topic to obtain knowledge corpus blocks. Each knowledge corpus describes the process of a penetration test and is recorded as ,in Indicates knowledge chunks, Indicates The task description of the knowledge corpus, Indicates A description of the solution steps for each knowledge chunk;

[0018] Use the text vector model to transform each knowledge corpus Mission Description Convert it into a semantic vector, store the semantic vector in a vector database, and create an index to complete the construction of the penetration testing solution knowledge base.

[0019] Preferably, the penetration testing tool library is constructed as follows:

[0020] Install the corresponding penetration testing tools on the attacking machine according to the requirements of the penetration test;

[0021] Provide a unified programming interface for each function of each penetration testing tool;

[0022] Provide a text description for each program interface, denoted as ,in, For the Programming interface, For the The interface calling format of a program interface, For the A functional description of the program interface, For the The parameter list description of the program interface, For the The interface returns a description of each program interface;

[0023] Use the text vector model to map each program interface Functional Description Convert it into a semantic vector, store the semantic vector in a vector database, and create an index to complete the construction of the penetration testing tool library.

[0024] Preferably, obtaining a list of penetration test steps according to the first prompt word includes:

[0025] Taking the first prompt word as input, the task planning agent returns a description of the solution steps;

[0026] Parse the solution step description to get a list of penetration test steps, denoted as ,in Indicates Description of the penetration testing steps, .

[0027] Preferably, the step of converting the tool output results of the current penetration test step and all penetration test steps before the current penetration test step into a second prompt word includes:

[0028] Get the first Tool output for each penetration testing step , and before The output results of the tools in the first penetration test step are combined to obtain the previous Execution trace of each penetration test step ;

[0029] The execution track With Description of the penetration testing steps Combined, we get Execution results of the penetration test step ;

[0030] Based on the description of the current penetration testing steps , No. Description of the penetration testing steps and Execution results of the penetration test step Transformed into the second prompt word.

[0031] Preferably, the operation type also includes execution stop. If the operation type is execution stop, the task scheduling agent Tool output for each penetration testing step Generates an error description and outputs the error description.

[0032] Preferably, the second prompt word includes a basis for determining the operation type, and the basis for determining the operation type is:

[0033] If the previous penetration test step is not the last penetration test step in the penetration test step list and the tool output result corresponding to the previous penetration test step includes execution is normal, the operation type of the current penetration test step is tool call;

[0034] If the previous penetration test step is the last penetration test step in the penetration test step list and the tool output result corresponding to the previous penetration test step includes execution normal, the operation type of the current penetration test step is execution summary;

[0035] If the tool output result corresponding to the previous penetration test step contains an execution exception, the operation type of the current penetration test step is execution stop.

[0036] Preferably, the combination of the current penetration test step, the previous penetration test step, and the execution result of the previous penetration test step is converted into a third prompt word, including:

[0037] Obtain the interface call format, function description and parameter list description corresponding to the retrieved program interface;

[0038] The third prompt word is converted according to the description of the current penetration test step, the description of the previous penetration test step, the execution result of the previous penetration test step, the interface call format, the function description and the parameter list description.

[0039] The invention provides a multi-agent-based large language model automatic penetration testing method, which has the following beneficial effects compared with the prior art:

[0040] (1) By integrating RAG (Retrieval-Augmented Generation) technology, the penetration testing process can be automatically executed without relying heavily on experienced professionals. This improvement not only improves the efficiency of penetration testing, but also significantly reduces labor costs.

[0041] (2) With the help of a multi-agent collaborative architecture and combined with the advantages of RAG technology, when dealing with complex penetration testing tasks, it can effectively reduce the hallucination problems that may occur in the decision-making process of large language models, thereby improving the completion rate and accuracy of penetration testing tasks. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] Figure 1 The present invention is a flowchart of a multi-agent-based large language model automatic penetration testing method. DETAILED DESCRIPTION

[0043] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0044] Unless otherwise defined, all technical and scientific terms used herein have the same meaning as those commonly understood by those skilled in the art of the present invention. The terms used in the specification of the present invention herein are only for the purpose of describing specific embodiments and are not intended to limit the present invention.

[0045] like Figure 1 As shown, this embodiment provides a multi-agent-based large language model automatic penetration testing method, including the following steps:

[0046] (1) Construction of penetration testing knowledge base: Penetration testing is highly professional, and general large language models (LLM) (especially open source general large language models) usually have little or no professional knowledge of penetration testing, resulting in serious hallucination problems. However, since the cost of retraining or fine-tuning large language models is too high (requiring a large number of training samples and high computing power), this embodiment uses RAG technology to assist the large language model in generating penetration testing task responses. Specifically, the description text of penetration testing expertise (including penetration testing solutions and penetration testing tools) is first stored in the knowledge base through collection, cleaning, indexing and other operations, and then injected into the large language model prompt.

[0047] (1-1) Construction of a penetration testing solution knowledge base: Since the general large language model lacks penetration testing expertise, the generated solutions are usually ineffective. Therefore, it is necessary to build a penetration testing solution knowledge base to alleviate this problem. The steps are as follows:

[0048] (1-1-1) Corpus collection: Based on data sources such as Wikipedia, CTFHub, and penetration testing papers, we obtain penetration testing professional knowledge corpus through web crawling.

[0049] (1-1-2) Knowledge Blocking: The knowledge corpus in the field of penetration testing is divided according to the topic to obtain a large number of knowledge corpus blocks. Each knowledge corpus describes the process of a penetration test and is recorded as ,in Indicates knowledge chunks, Indicates The task description of the knowledge corpus (i.e. the expected goal description of the penetration test), Indicates A description of the solution steps for a knowledge corpus (i.e. a description of the specific steps of the penetration test).

[0050] (1-1-3) Knowledge storage: First, use a text vector model (such as the embedding model) to store each knowledge corpus. Mission Description The semantic vectors are then stored in a vector database and indexed.

[0051] (1-2) Construction of penetration testing tool library: Penetration testing requires the use of various penetration testing tools to complete. In order to facilitate the large language model to automatically retrieve and execute the corresponding penetration testing tools, it is necessary to build a penetration testing tool library in advance. The steps are as follows:

[0052] (1-2-1) Tool installation: Install corresponding tools (such as metasploit tool and nmap tool) on the attacking machine according to the requirements of the penetration test.

[0053] (1-2-2) Programming interface for tool calls: Different tools have different calling methods (some tools provide API development interfaces, some tools need to be accessed through command lines, and some tools may contain multiple functions). To this end, different calling methods are hidden by writing programs, and a unified program interface is provided for each function. For example, shell commands are called through the SSH protocol to process.

[0054] (1-2-3) Program interface text description: Provide a series of text descriptions for each program interface, recorded as ,in, For the Programming interface, For the The interface calling format of a program interface, For the A functional description of the program interface, For the A parameter list description of the program interface (including the meaning and type of each parameter), For the The interface returns a description of the program interface (i.e. as a tool output result).

[0055] (1-2-4) Program interface storage: First, use a text vector model (such as the emmbeding model) to store each program interface. Functional Description The semantic vectors are then stored in a vector database and indexed.

[0056] (2) Penetration testing based on large-model intelligent agents: By designing a variety of large-model intelligent agents, a series of decision-making and operational tasks such as knowledge retrieval, task planning, task scheduling, tool calling, and task summary are automatically completed in a division of labor, thus realizing the automation of the penetration testing process.

[0057] (2-1) Agent definition: The following five large-scale model agents are defined:

[0058] (2-1-1) Knowledge retrieval agent: used to perform retrieval operations in RAG technology. Specifically, given a task description entered by the user First, use the text vector model to describe the task Convert it into a semantic vector and retrieve the most relevant knowledge chunks from the penetration testing solution knowledge base based on its index , which is used to subsequently enhance the prompt words input into the large language model.

[0059] (2-1-2) Task Planning Agent: It is used to generate the steps of the penetration test task. Specifically, the task description entered by the user is , the knowledge corpus obtained by searching The description of the constraints is organized into the first prompt word, and the solution step description for completing the input task is generated according to the first prompt word, and a list of penetration test steps is obtained, which is recorded as ,in Indicates Description of the penetration testing steps, .

[0060] (2-1-3) Task Scheduling Agent: It is used to generate the execution plan for the single-step penetration test step. Specifically, the description of the current penetration test step Execution trace of all previous penetration testing steps and the execution environment description are organized into a second prompt word, where Indicates The tool output result corresponding to the first penetration test step is used to generate the operation type of the current penetration test step according to the second prompt word, which is recorded as The operation types of this embodiment include tool calling, execution summary and execution stop. For tool call, start the tool call agent; if To perform the summary, the task summary agent is started; if If the execution stops, it means that an error occurs in the penetration test automation process, and manual processing is required.

[0061] (2-1-4) Tool Calling Agent: Used to call the penetration testing tool to perform the expected operation. Specifically, it receives the description of the current penetration testing step from the task scheduling agent. , Description of the previous penetration test step , the execution result of the previous penetration test step As input, the description of the current penetration testing step Convert it into a semantic vector and retrieve the most relevant program interface based on the index of the penetration testing tool library , the description of the current penetration test step , Description of the previous penetration test step , the execution result of the previous penetration test step , and program interfaces Corresponding interface call format , Parameter list description , Functional Description and the execution environment description are organized into a third prompt word, and a tool calling code for calling the corresponding tool is generated according to the third prompt word, the tool calling code is executed to call the penetration testing tool to obtain the returned tool output result, and the tool output result is returned to the task scheduling intelligent body.

[0062] (2-1-5) Task summary agent: It is used to generate a summary of the task after all steps are completed. Specifically, the input task description and execution traces of all penetration testing steps The fourth prompt word is organized, and a task summary description is obtained according to the fourth prompt word as an automatic penetration test result.

[0063] (2-2) Large-model agent collaboration: receiving task descriptions input by users Finally, the execution process of the multi-agent-based penetration test task is as follows:

[0064] (2-2-1) Call the knowledge retrieval agent according to the task description Retrieve knowledge blocks from the penetration testing solution knowledge base .

[0065] (2-2-2) Call the task planning agent and describe the task according to the first prompt word template , the knowledge corpus obtained by searching Organize into the first prompt word, and the task planning agent returns the solution step description according to the first prompt word, and parses the solution step description to obtain the penetration test step list The first prompt word template provided in this embodiment is as follows:

[0066] As a penetration test task planning assistant, your task is to describe the user's penetration test goals , reference penetration testing solutions , generate a description of the solution steps to accomplish the overall goal.

[0067] ;

[0068] .

[0069] (2-2-3) List of penetration testing steps Each step of the penetration test , calling the task scheduling agent, and according to the second prompt word template, the description of the current penetration test step Execution trace of all previous penetration testing steps Organize into the second prompt word, specifically first obtain the first prompt word returned by the tool calling agent Tool output for each penetration testing step , and before The output results of the tools in the first penetration test step are combined to obtain the previous Execution trace of each penetration test step , and then combine to get the execution result of the previous penetration test step , according to the second prompt word template, it is converted into the second prompt word, and the task scheduling agent returns the operation type according to the second prompt word If the operation type For tool call, go to step (2-2-4); if To perform the summary, go to step (2-2-5); if To stop execution, the task scheduling agent outputs the results according to the tool output of the previous penetration test step Generate an error description and go to step (2-2-6). The description of the current penetration test step hour, This embodiment provides the second prompt word template as follows:

[0070] As a penetration test step execution assistant, your task is to follow the description of the current penetration test step , Description of the previous step , the execution result of the previous step , determine the operation type of the current step.

[0071] The operation types include: tool call, execution summary, and execution stop.

[0072] The judgment basis is: if the previous step is not the last step and is executed normally, the operation type is tool call; if the previous step is the last step and is executed normally, the operation type is execution summary; if the previous step is executed abnormally, the operation type is execution stop.

[0073] ;

[0074] ;

[0075] .

[0076] (2-2-4) Start the tool to call the agent, and write the description of the current penetration test step according to the third prompt word template , Description of the previous penetration test step , the execution result of the previous penetration test step And the program interface Corresponding interface call format , Parameter list description , Functional Description The tool calling agent returns the tool calling code according to the third prompt word, executes the tool calling code, obtains the tool output result returned by the penetration test tool, and returns the tool output result to step (2-2-3), and executes the next penetration test step. The third prompt word template provided in this embodiment is as follows:

[0077] As a program interface call assistant, your task is to follow the description of the current penetration test steps , description of the previous step , the execution result of the previous step , and program interfaces Corresponding interface call format , Parameter list description and functional description Generate tool calling code.

[0078] ;

[0079] ;

[0080] ;

[0081] .

[0082] (2-2-5) Start the task summary agent and summarize the task description according to the fourth prompt word template , complete step execution trajectory , the execution results of all steps are actually organized into a fourth prompt word, and a task summary description is output to the user as the automatic penetration test result according to the fourth prompt word. This embodiment provides the fourth prompt word template as follows:

[0083] As a penetration test task management assistant, your task is to describe And the complete penetration test step execution results , and give a summary of the execution process of this penetration test.

[0084] ;

[0085] .

[0086] (2-2-6) Output the error description and request manual processing.

[0087] Comparative experiment: In order to verify the effectiveness and superiority of the large model agent integrating RAG technology in penetration testing, this embodiment conducted a comparative experiment. The experiment aims to compare the differences between the large language model automated penetration testing solution (penetration testing without RAG support) and the large language model automated testing solution enhanced by RAG technology (penetration testing with RAG support), and record the test results, especially focusing on the performance in terms of test efficiency, accuracy and success rate.

[0088] Experimental environment: Test platform: a unified virtual environment to ensure the consistency of experimental conditions. Test objectives: select the same target system for penetration testing. Test tools: use the same set of automated penetration testing tools, the only variable is whether there is support for RAG technology.

[0089] Experimental steps:

[0090] Penetration testing without RAG support: Start from step (2-2-2) but do not add the retrieved knowledge chunks , that is, do not use RAG technology, automatically execute until the task is completed or an error description is output. Repeat the test multiple times to reduce the impact of accidental factors, and record the corresponding results.

[0091] Penetration testing with RAG support: In the same test environment, use an automated penetration testing tool that incorporates RAG technology. Repeat the penetration test in step 2 until the task is completed or an error description is output. Repeat the test multiple times to reduce the impact of accidental factors and record the corresponding results.

[0092] Experimental results:

[0093] Penetration test results without RAG support: Test duration: 1 hour; Penetration tasks completed: 0 times; Number of false alarms: 5 times; Test success rate: 0% (defined as the proportion of successful penetration to total penetration tasks).

[0094] Penetration test results with RAG support: Test duration: 1 hour; Penetration tasks completed: 5 times; Number of false alarms: 0 times; Test success rate: 100% (defined as the proportion of successful penetrations to total penetration tasks).

[0095] This experiment shows that the application of large model agents integrating RAG technology in automated penetration testing not only improves the efficiency of the test, but also greatly improves the accuracy and success rate of the test. The application of the present invention helps enterprises better identify and repair potential security threats in actual operations and ensure the security of information systems.

[0096] The technical features of the above-described embodiments may be arbitrarily combined. To make the description concise, not all possible combinations of the technical features in the above-described embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.

[0097] The above-mentioned embodiments only express several implementation modes of the present invention, and the description thereof is relatively specific and detailed, but it cannot be understood as limiting the scope of the invention. It should be pointed out that, for a person of ordinary skill in the art, several modifications and improvements can be made without departing from the concept of the present invention, and these all belong to the protection scope of the present invention. Therefore, the protection scope of the present invention shall be subject to the attached claims.

Claims

1. A multi-agent based large language model automatic penetration testing method, characterized in that: The multi-agent-based large language model automatic penetration testing method comprises: Build a penetration testing solution knowledge base and a penetration testing tool library; Establish five large-scale model agents, including knowledge retrieval agent, task planning agent, task scheduling agent, tool calling agent and task summary agent; Receive an input task description, call the knowledge retrieval agent, and retrieve a knowledge corpus in the penetration testing solution knowledge base according to the task description; Calling the task planning agent to convert the task description and the knowledge corpus into a first prompt word, and obtaining a penetration test step list according to the first prompt word; For each penetration test step in the penetration test step list, call the task scheduling agent, convert the tool output results of the current penetration test step and all penetration test steps before the current penetration test step into the second prompt word, and obtain the operation type corresponding to the current penetration test step according to the second prompt word. If the operation type is tool call, call the tool call agent; if the operation type is execution summary, call the task summary agent; The tool calling agent retrieves the corresponding program interface in the penetration testing tool library according to the current penetration testing step, combines the current penetration testing step, the previous penetration testing step, and the execution result of the previous penetration testing step into a third prompt word, obtains the tool output result corresponding to the current penetration testing step according to the third prompt word, and inputs the tool output result into the task scheduling agent; The task summary agent converts the task description and the execution results of all penetration test steps into a fourth prompt word, and outputs the automatic penetration test result according to the fourth prompt word.

2. The multi-agent-based large language model automatic penetration testing method according to claim 1 is characterized in that: The penetration testing solution knowledge base is constructed as follows: Acquire penetration testing professional knowledge corpus through web crawling; The knowledge corpus in the field of penetration testing is divided according to the topic to obtain knowledge corpus blocks. Each knowledge corpus describes the process of a penetration test and is recorded as ,in Indicates knowledge chunks, Indicates The task description of the knowledge corpus, Indicates A description of the solution steps for each knowledge chunk; Use the text vector model to transform each knowledge corpus Mission Description Convert it into a semantic vector, store the semantic vector in a vector database, and create an index to complete the construction of the penetration testing solution knowledge base.

3. The multi-agent-based large language model automatic penetration testing method according to claim 1 is characterized in that: The construction process of the penetration testing tool library is as follows: Install the corresponding penetration testing tools on the attacking machine according to the requirements of the penetration test; Provide a unified programming interface for each function of each penetration testing tool; Provide a text description for each program interface, denoted as ,in, For the Programming interface, For the The interface calling format of a program interface, For the A functional description of the program interface, For the The parameter list description of the program interface, For the The interface returns a description of each program interface; Use the text vector model to map each program interface Functional description Convert it into a semantic vector, store the semantic vector in a vector database, and create an index to complete the construction of the penetration testing tool library.

4. The multi-agent-based large language model automatic penetration testing method according to claim 1 is characterized in that: The step of obtaining a penetration test step list according to the first prompt word includes: Taking the first prompt word as input, the task planning agent returns a description of the solution steps; Parse the solution step description to get a list of penetration test steps, denoted as ,in Indicates Description of the penetration testing steps, .

5. The multi-agent-based large language model automatic penetration testing method according to claim 1 is characterized in that: The step of converting the tool output results of the current penetration test step and all penetration test steps before the current penetration test step into a second prompt word includes: Get the first Tool output for each penetration testing step , and before The output results of the tools in the first penetration test step are combined to obtain the previous Execution trace of each penetration test step ; The execution track With Description of the penetration testing steps Combined, we get Execution results of the penetration test step ; Based on the description of the current penetration testing steps , No. Description of the penetration testing steps and Execution results of the penetration test step Transformed into the second prompt word.

6. The multi-agent-based large language model automatic penetration testing method according to claim 1 is characterized in that: The operation type also includes execution stop. If the operation type is execution stop, the task scheduling agent Tool output for each penetration testing step Generates an error description and outputs the error description.

7. The multi-agent-based large language model automatic penetration testing method according to claim 6 is characterized in that: The second prompt word includes a basis for determining the operation type, and the basis for determining the operation type is: If the previous penetration test step is not the last penetration test step in the penetration test step list and the tool output result corresponding to the previous penetration test step includes execution is normal, the operation type of the current penetration test step is tool call; If the previous penetration test step is the last penetration test step in the penetration test step list and the tool output result corresponding to the previous penetration test step includes execution normal, the operation type of the current penetration test step is execution summary; If the tool output result corresponding to the previous penetration test step contains an execution exception, the operation type of the current penetration test step is execution stop.

8. The multi-agent-based large language model automatic penetration testing method according to claim 1 is characterized in that: The combination of the current penetration test step, the previous penetration test step, and the execution result of the previous penetration test step is converted into a third prompt word, including: Obtain the interface call format, function description and parameter list description corresponding to the retrieved program interface; The third prompt word is converted according to the description of the current penetration test step, the description of the previous penetration test step, the execution result of the previous penetration test step, the interface call format, the function description and the parameter list description.

Images