Reinforcement learning based attack agent penetration testing method and system

Abstract

The application discloses a kind of based on reinforcement learning's attack intelligent agent penetration testing method and system, based on reinforcement learning's attack intelligent agent penetration testing method includes: according to DQN algorithm, Actor-Critic mechanism and preset priori knowledge framework constructs training thread model, according to A3C algorithm, multiple training thread models are synchronously trained, and decision model is obtained;Network environment information is obtained;According to network environment information and decision model, determine the current tactical phase, and then obtain the optimal attack action signal;According to the optimal attack action signal, initiate simulated attack to target system, and record test results;According to test results, obtain the first current hidden state and message vector of decision model, obtain policy tendency information according to the first current hidden state and message vector, and update other decision models according to policy tendency information.The application can improve the intelligent degree of penetration testing and the efficiency of penetration testing in large network, and can be widely applied in automatic penetration field.

Description

Technical Field

[0001] This invention relates to the field of automated penetration testing, and in particular to a penetration testing method and system for attack agents based on reinforcement learning. Background Technology

[0002] Penetration testing is an authorized, simulated hacker attack method that assesses the weaknesses and vulnerabilities of a target system by mimicking the behavior of an attacker. Automated penetration testing is an extension of penetration testing that uses automated tools and scripts to execute attacks. With the development of artificial intelligence, some automated penetration testing tools have begun to utilize reinforcement learning techniques to improve the efficiency and success rate of attacks. By allowing the tools to continuously learn and optimize attack strategies, they can better adapt to changes in the target system. Current automated penetration testing methods can be broadly classified into the following three categories based on their level of intelligence:

[0003] 1. Mapping the results of vulnerability scanners to corresponding penetration testing tools. This method often integrates modules such as vulnerability scanning, penetration attack, and payload selection into a framework to achieve a basic automated penetration testing process. It can only achieve automated penetration testing with a low level of intelligence, relies heavily on behavioral rules set by human experts, and does not have intelligent capabilities.

[0004] Second, the attack and defense process is modeled as an attack graph and decision tree, and the automated penetration testing process is transformed into a penetration path planning problem for solution. By displaying the possible attack sequence and effects of the attacker, the attack path can be selected by using the known vulnerabilities and the relationships between them. This method achieves a certain degree of intelligent path planning, but it is not suitable for large networks and cannot avoid the dependence on human decision-making and operation.

[0005] III. Artificial Intelligence-Based Policy Transfer Methods: Most existing AI-based intelligent penetration schemes focus on the decision-making effect of a single agent in a small network. When the target network is a large network, the action space and state space in reinforcement learning will grow exponentially. Traditional single-agent automatic schemes face problems such as training convergence difficulties and low efficiency. Summary of the Invention

[0006] To address the aforementioned technical problems, the present invention aims to provide a penetration testing method and system for attack agents based on reinforcement learning, which can improve the intelligence level of the system and effectively enhance the efficiency of penetration testing in large-scale network scenarios.

[0007] The technical solution adopted in one aspect of this invention is:

[0008] A penetration testing method for attack agents based on reinforcement learning includes the following steps:

[0009] Multiple training thread models are constructed based on the DQN algorithm, the Actor-Critic mechanism, and a pre-defined prior knowledge framework. The multiple training thread models are then trained synchronously using the A3C algorithm to obtain multiple decision models.

[0010] Obtain network environment information;

[0011] The current tactical stage is determined based on the network environment information and the decision model, and then the optimal attack action signal is obtained based on the current tactical stage.

[0012] A simulated attack was launched on the target system based on the optimal attack action signal, and the test results were recorded.

[0013] Based on the test results, the first current hidden state and message vector of the corresponding decision model are obtained through the message encoder. Then, based on the first current hidden state and the message vector, the policy tendency information is obtained through the message decoder and message aggregator. The other decision models are updated based on the policy tendency information.

[0014] Furthermore, the prior knowledge framework includes:

[0015] The penetration testing process is divided into multiple tactical phases, and multiple attack actions are defined according to each tactical phase, resulting in a prior knowledge framework.

[0016] Furthermore, the step of constructing multiple training thread models based on the DQN algorithm, the Actor-Critic mechanism, and a pre-defined prior knowledge framework specifically includes:

[0017] The reinforcement learning network that applies the attention mechanism is used as the value network Critic, and the pre-defined deep neural network is used as the policy network Actor.

[0018] The attack action is evaluated using the value network Critic based on the prior knowledge framework to obtain the attack action reward value;

[0019] The policy network Actor is updated based on the attack action reward value to obtain multiple training thread models.

[0020] Furthermore, the step of synchronously training multiple training thread models according to the A3C algorithm to obtain multiple decision models specifically includes:

[0021] Simultaneously, multiple threads are started, and multiple network environment sample information is obtained;

[0022] Decisions are made based on the network environment sample information described above, and corresponding reward signals are obtained;

[0023] Update the gradient parameters of the corresponding training thread model according to each of the reward signals;

[0024] The common model is updated based on the gradient parameters described above, and the number of updates is recorded.

[0025] When the number of updates exceeds the preset number of updates, the training thread models are trained according to the common model to obtain the decision model.

[0026] Furthermore, the step of launching a simulated attack on the target system based on the optimal attack action signal and recording the test results specifically includes:

[0027] Based on the optimal attack action signal, a corresponding attack payload is generated by the attack payload generator.

[0028] Based on the attack payload, a simulated attack is launched on the target system using a penetration testing tool, and the test results are recorded.

[0029] Furthermore, the step of obtaining the first current hidden state and message vector of the decision model corresponding to the test results through the message encoder specifically includes:

[0030] When the test result reaches the preset result, the hidden state of the corresponding decision model at the previous moment is obtained;

[0031] The previous hidden state is input into the message encoder, and the first current hidden state and the message vector are output.

[0032] Furthermore, the step of obtaining policy bias information based on the first current hidden state and the message vector through a message decoder and a message aggregator, and updating other decision models based on the policy bias information, specifically includes:

[0033] The first current hidden state and the message vector are input into the message decoder, and the second current hidden state and attention weights are output.

[0034] The second current hidden state and the attention weight are fused to obtain the policy tendency information;

[0035] The policy preference information is input into other decision models to update them.

[0036] Another aspect of the technical solution adopted by the present invention is:

[0037] A penetration testing system for attack agents based on reinforcement learning, comprising:

[0038] The training module is used to construct multiple training thread models based on the DQN algorithm, the Actor-Critic mechanism, and a preset prior knowledge framework, and to synchronously train the multiple training thread models according to the A3C algorithm to obtain multiple decision models.

[0039] The sensing module is used to acquire network environment information;

[0040] The decision module is used to determine the current tactical stage based on the network environment information and the decision model, and then obtain the optimal attack action signal based on the current tactical stage.

[0041] The control module is used to launch a simulated attack on the target system according to the optimal attack action signal and record the test results;

[0042] The communication module is used to obtain the first current hidden state and message vector of the corresponding decision model through the message encoder based on the test results, and then obtain policy tendency information through the message decoder and message aggregator based on the first current hidden state and the message vector, and update other decision models based on the policy tendency information.

[0043] Another aspect of the technical solution adopted by the present invention is:

[0044] A computer device, comprising:

[0045] At least one processor;

[0046] At least one memory for storing at least one program;

[0047] When the at least one program is executed by the at least one processor, the at least one processor implements the aforementioned attack agent penetration testing method.

[0048] Another aspect of the technical solution adopted by the present invention is:

[0049] A computer-readable storage medium storing a processor-executable program, characterized in that the processor-executable program, when executed by the processor, is used to implement the aforementioned attack agent penetration testing method.

[0050] The beneficial effects of this invention are as follows: The penetration testing method and system based on reinforcement learning of this invention, on the one hand, can autonomously observe the network environment and make decisions by using a decision model trained synchronously based on the DQN algorithm, the Actor-Critic mechanism, and the A3C algorithm, as well as a preset prior knowledge framework, thereby improving the intelligence level and penetration testing efficiency of the system; on the other hand, by aggregating the hidden states and message vectors of the decision model to obtain policy bias information, and transmitting the policy bias information to other decision models in the same network environment for model updates, multi-agent collaborative penetration can be achieved in large-scale network scenarios. Attached Figure Description

[0051] Figure 1 A flowchart illustrating the steps of the penetration testing method for attacking intelligent agents provided in this embodiment of the invention;

[0052] Figure 2 An architecture diagram of the training thread model provided in an embodiment of the present invention;

[0053] Figure 3 This is a schematic diagram of the training process of a decision model based on the A3C algorithm provided in an embodiment of the present invention;

[0054] Figure 4 This is a diagram illustrating the working mechanism of communication between multiple attack agents provided in an embodiment of the present invention.

[0055] Figure 5 This is a schematic diagram of the structure of the attack agent penetration testing system provided in an embodiment of the present invention;

[0056] Figure 6 This is a schematic diagram of the working framework of the penetration testing process provided in an embodiment of the present invention;

[0057] Figure 7 A schematic diagram illustrating the workflow of a penetration testing process provided in an embodiment of the present invention;

[0058] Figure 8 This is a schematic diagram of the structure of a computer device provided in an embodiment of the present invention. Detailed Implementation

[0059] Penetration testing is an authorized, simulated hacker attack method that assesses the weaknesses and vulnerabilities of a target system by mimicking attacker behavior. Automated penetration testing is an extension of penetration testing that uses automated tools and scripts to execute attacks. With the development of artificial intelligence, some automated penetration testing tools have begun to utilize reinforcement learning techniques to improve the efficiency and success rate of attacks. By allowing the tools to continuously learn and optimize attack strategies, they can better adapt to changes in the target system. Existing automated penetration testing methods can be broadly classified into three categories based on their level of intelligence. The most basic approach to automation maps vulnerability scanner results to corresponding penetration testing tools. This method often integrates vulnerability scanning, penetration attacks, and payload selection into a single framework to achieve a basic level of automated penetration testing. It only achieves a low level of intelligence, heavily relying on rules defined by human experts and lacking true intelligence. The second type of approach models the attack and defense process as attack graphs and decision trees, transforming the automated penetration testing process into a penetration path planning problem. This type of method is primarily based on the attack graph proposed by Phillips et al. in 1998. By displaying the possible attack sequence and effects, attackers can leverage known vulnerabilities and their relationships to select attack paths. This approach achieves a... While offering some degree of intelligent path planning, it is not suitable for large networks and cannot avoid reliance on human decision-making and operation. The third type of method is AI-based PT methods. Based on the AI ​​technology used in PT, AI-based intelligent PT methods can be mainly divided into two categories: PT based on traditional reinforcement learning (RL) and PT based on deep reinforcement learning. The latter combines deep learning and reinforcement learning, and under sufficient training conditions, it fits the value function more accurately. Most existing AI-based intelligent penetration schemes focus on the decision-making effect of a single agent in small networks. When the target network is a large network, the action space and state space in reinforcement learning will grow exponentially. Traditional single-agent automatic schemes face difficulties in training convergence and the explosion of state and action spaces.

[0060] To address this, this invention proposes a penetration testing method for attack agents based on reinforcement learning. On one hand, by using a decision model trained synchronously based on the DQN algorithm, the Actor-Critic mechanism, and the A3C algorithm, along with a pre-defined prior knowledge framework, the system can autonomously observe the network environment and make decisions, thereby improving the intelligence level and penetration testing efficiency of the system. On the other hand, by aggregating the hidden states and message vectors of the decision model to obtain policy bias information, and transmitting this policy bias information to other decision models in the same network environment for model updates, multi-agent collaborative penetration can be achieved in large-scale network scenarios.

[0061] This invention proposes a penetration testing method for attack agents based on reinforcement learning, such as... Figure 1 As shown, the method includes steps S101 to S105:

[0062] S101. Construct multiple training thread models based on the DQN algorithm, the Actor-Critic mechanism, and the preset prior knowledge framework. Simultaneously train the multiple training thread models using the A3C algorithm to obtain multiple decision models.

[0063] Specifically, by using a pre-defined prior knowledge framework, the penetration steps and specific scenario behavioral logic are subdivided and generalized, thereby helping the decision model build basic penetration decision logic to quickly initialize the model and provide a stable basic framework for the subsequent decision logic of the decision model. Multiple decision models are obtained through synchronous training using the DQN algorithm, the Actor-Critic mechanism, and the A3C algorithm.

[0064] As an optional implementation, the prior knowledge framework includes:

[0065] The penetration testing process is divided into multiple tactical phases, and multiple attack actions are defined according to each tactical phase, resulting in a prior knowledge framework.

[0066] Specifically, this embodiment of the invention adopts the Attack Tactics Standard (ATT&CK) as the basic framework for prior knowledge delineation, dividing attack actions into 14 tactical phases and hundreds of specific technical actions, such as Collection, Privilege Escalation, Defense Evasion, Discovery, Credential Access, Lateral Movement, Execution, Persistence, Impact, Command and Control, Exfiltration, Impact, Preparation, and Initial Access. Under these 14 phases, further subdivisions of attack actions that can be performed in different phases are provided. For example, in the tactical phase "Persistence," tactical actions such as registry startup items and service installation are included, describing the actions an attacker might take at this phase. Addressing the difficulty of agent training and fitting in automated penetration testing tasks, this embodiment of the invention rapidly initializes the model's behavioral framework through a prior knowledge framework, thereby effectively shortening training time and improving environmental adaptability.

[0067] As an optional implementation, the step of constructing multiple training thread models based on the DQN algorithm, the Actor-Critic mechanism, and a pre-defined prior knowledge framework can be divided into the following steps S1011 to S1013:

[0068] S1011. Use the reinforcement learning network with attention mechanism as the value network Critic and the pre-defined deep neural network as the policy network Actor.

[0069] Specifically, such as Figure 2 As shown, this embodiment of the invention uses Actor-Critic as the main framework for training the thread model, employs a DQN network with an attention mechanism as the value network Critic, and defines a new deep neural network as the policy network Actor. The attention mechanism in the value network Critic guides the network's convolutional layers to focus on the information most valuable for decision-making, thereby reducing the dimensionality of the dataset and the large amount of computation on the original data, and thus supporting the decision model to make the decision with the highest long-term value.

[0070] S1012. Based on the prior knowledge framework, the attack action is evaluated through the value network Critic to obtain the attack action reward value;

[0071] Specifically, in the reinforcement learning reward setting process, this embodiment of the invention sets a reward value for each attack action and tactic. Successfully executed attack actions and tactics can receive positive rewards, while failed attempts may be penalized. Simultaneously, considering the resource consumption and time factors of attack decision-making, different reward values ​​are set for different attack actions and tactics, with the maximum reward set for successfully gaining control of the target system to encourage the model to continuously attempt towards the final goal. The value network Critic uses a value function and a dominance function to evaluate the attack actions of the policy network Actor and feeds back the attack action reward value to the policy network Actor.

[0072] S1013. Update the policy network Actor according to the attack action reward value to obtain the training thread model.

[0073] Specifically, the policy network (Actor) updates its policy based on the attack action reward value fed back by the value network (Critic), dynamically changing the policy to adapt to the continuous network environment. In each interaction with the environment, the model obtains corresponding reward values ​​based on the tactical stage reached by the agent and the attack action executed. With the aim of maximizing long-term rewards, the model continuously updates the value network (Critic) and the policy network (Actor) during interactions with the environment, achieving accurate simulation of state-value pairs and optimization of the policy. Addressing the issue of low intelligence in existing automated penetration testing schemes, this embodiment of the invention uses the DQN algorithm combined with the Actor-Critic mechanism to perform intelligent decision-making based on the characteristics of different network environments.

[0074] Reference Figure 3 As an optional implementation, the step of simultaneously training multiple training thread models according to the A3C algorithm to obtain multiple decision models can be further divided into the following steps S1014 and S1018:

[0075] S1014. Simultaneously start multiple threads and obtain multiple network environment sample information;

[0076] Specifically, during the training process, three threads are simultaneously started to train three training thread models (all with the same architecture, i.e., the A3C algorithm). The system uses scanning behavior to collect information in a small range to identify the current network composition, including host PCs, servers, routers, and security devices. The overall state perceived by the system includes the state of the attacking agent and the state of the network environment. The state of the attacking agent includes its functions and the number of Proof-of-Concept (POCs). The state of the network environment includes its components, i.e., network devices, services, and vulnerabilities. Then, the inputs from different network sensors and processes are fused and prepared to be input into the deep neural network of the training thread model. Each network sensor information captures the original state information of the environment in the form of raw vectors. Each information group performs information aggregation processing according to the characteristics of different environmental information and feeds it into the deep neural network of the decision-making stage.

[0077] S1015. Make decisions based on the sample information of each network environment and obtain the corresponding reward signal;

[0078] Specifically, each training thread model makes decisions based on network environment sample information and receives its corresponding reward signal based on the effect of the decision.

[0079] S1016. Update the gradient parameters of the corresponding training thread model according to each reward signal;

[0080] Specifically, through continuous interaction with the training environment, the three training thread models will continuously optimize the gradient parameters in the network, thereby continuously improving the efficiency and accuracy of decision-making.

[0081] S1017. Update the common model according to each gradient parameter and record the number of updates;

[0082] Specifically, after the three threads have been trained for a period of time, the gradient parameters of the common model are updated using their respective gradient parameters, and the number of updates is recorded.

[0083] S1018. When the number of updates exceeds the preset number of updates, train the models of each training thread according to the common model to obtain multiple decision models.

[0084] Specifically, when the number of updates to the common model reaches a set value, each training thread will copy the gradient parameters from the common model and perform new training to obtain the decision model. This ensures the continuity and sufficiency of the gradient parameters in the common model during training and reduces the probability of information conflicts caused by multi-threaded training. Updating the common model after synchronous multi-threaded training means that the training time is shortened significantly, improving training efficiency. Addressing the problem of low efficiency and difficulty in agent training in large-scale network scenarios of existing automatic penetration testing schemes, this embodiment of the invention uses the A3C training scheme in automatic penetration testing scenarios to accelerate the training process of the decision model through asynchronous training.

[0085] S102. Obtain network environment information;

[0086] Specifically, scanners are used to perceive the current network composition and monitor the composition information of the network environment, such as host PCs, servers, routers, and protection devices, thereby capturing key network data to obtain network environment information. This network environment information is then transformed into feature vectors and aggregated.

[0087] S103. Determine the current tactical stage based on network environment information and decision-making model, and then obtain the optimal attack action signal based on the current tactical stage;

[0088] Specifically, the decision model receives feature vectors from the aggregated network environment information, analyzes and processes these data through deep neural networks, finds hidden representations, predicts the best attack action, and simulates an attack on the target system. The decision model will judge the current tactical stage of the attack and, based on the tactical actions that can be executed at the current tactical stage in the prior knowledge framework, decide the tactical action with the highest value at the current stage, and transmit the optimal attack action signal to the next module of the system to carry out the penetration behavior.

[0089] S104. Launch a simulated attack on the target system based on the optimal attack action signal and record the test results;

[0090] As an optional implementation, the step of launching a simulated attack on the target system based on the optimal attack action signal and recording the test results can be further divided into the following steps S1041 and S1042:

[0091] S1041. Generate the corresponding attack payload through the attack payload generator based on the optimal attack action signal;

[0092] Specifically, a payload refers to the attack path taken by a penetration testing agent when it interacts with a target system to launch an attack. A payload generator is a tool used to create and generate actual attack packets or code to attack a target system. In other words, the payload generator will generate code that directly interacts with the target system and find effective payloads.

[0093] S1042. Based on the attack payload, launch a simulated attack on the target system using a penetration testing tool and record the test results.

[0094] Specifically, penetration testing tools refer to tools that simulate attacks on target systems to detect their vulnerabilities and security. These tools can execute various vulnerability exploitation and attack techniques, such as vulnerability exploitation tools, password cracking tools, malware tools, Trojans and backdoor tools, network sniffing, DNS spoofing, wireless attacks, social engineering attacks, and so on, to simulate diverse attack scenarios and conduct simulated attacks on target systems.

[0095] As mentioned earlier, when the attacking agent receives the optimal attack action signal, it first activates the attack payload generator to generate the corresponding attack payload, which is a command that directly interacts with the target system. The command will call the corresponding penetration testing tools to launch a simulated attack on the target system according to the interaction stage with the target system and the type of different vulnerabilities. After the attack, new network environment information is obtained again. The environmental feedback of the previous attack action is obtained from the new network environment information, and this feedback is delivered to the decision module to make a new decision.

[0096] S105. Based on the test results, obtain the first current hidden state and message vector of the corresponding decision model through the message encoder, and then obtain the policy tendency information through the message decoder and message aggregator based on the first current hidden state and message vector, and update other decision models based on the policy tendency information.

[0097] Specifically, after a certain number of deliveries are made to the system and environment, the embodiments of the present invention encrypt and transmit the network environment features learned by the current attack agent to other attack agents in the same network environment, thereby accelerating the perception and adaptation of other attack agents to the current network defense style and providing the possibility for multi-agent collaboration in large-scale network systems.

[0098] Reference Figure 4Furthermore, as an optional implementation, the step of obtaining the first current hidden state and message vector of the decision model through the message encoder based on the test results can be further divided into the following steps S1051 and S1052:

[0099] S1051. When the test result reaches the preset result, obtain the hidden state of the corresponding decision model at the previous moment.

[0100] S1052. Input the hidden state of the previous time step into the message encoder and output the first current hidden state and message vector;

[0101] Specifically, after the system interacts with the network environment a certain number of times, the system will learn the characteristics of the current network environment. The current attacking agent takes its local observations and the hidden state of the previous time step as input through the message encoder, and outputs the hidden state of the current time step and the message vector. The message vector is a continuous vector that contains the local information and historical information of the system. The message encoder is a recurrent neural network (RNN) with memory capabilities.

[0102] Reference Figure 4 As an optional implementation, the step of obtaining policy bias information through a message decoder and a message aggregator based on the first current hidden state and message vector, and updating other decision models based on the policy bias information, can be further divided into the following steps S1053 and S1055:

[0103] S1053. Input the first current hidden state and message vector into the message decoder and output the second current hidden state and attention weights;

[0104] Specifically, other attack agents receive encrypted message vectors and hidden states from the current attack agent. They use the received encrypted message vectors and hidden states as input through a message decoder, and then output the hidden state and attention weights at the current moment. The attention weights are normalized vectors that represent the importance of the receiving attack agent to the messages of other attack agents. The message decoder is also composed of a recurrent neural network (RNN).

[0105] S1054. Fuse the second current hidden state and attention weights to obtain policy bias information;

[0106] S1055. Input the strategy preference information into other decision-making models and update the other decision-making models.

[0107] Specifically, the outputs of the message decoder are fused by the message aggregator to obtain policy bias information, which is then input into the decision models of other attack agents to update and adjust other decision models in the same network environment.

[0108] Reference Figure 5 This invention also provides a penetration testing system for attack agents based on reinforcement learning, comprising:

[0109] The training module is used to construct multiple training thread models based on the DQN algorithm, the Actor-Critic mechanism, and a preset prior knowledge framework. The multiple training thread models are trained synchronously according to the A3C algorithm to obtain multiple decision models.

[0110] The sensing module is used to acquire network environment information;

[0111] The decision-making module is used to determine the current tactical stage based on network environment information and decision-making models, and then obtain the optimal attack action signal based on the current tactical stage.

[0112] The control module is used to launch simulated attacks on the target system based on the optimal attack action signal and record the test results;

[0113] The communication module is used to obtain the first current hidden state and message vector of the corresponding decision model through the message encoder based on the test results, and then obtain policy tendency information through the message decoder and message aggregator based on the first current hidden state and message vector, and update other decision models based on the policy tendency information.

[0114] The attack agent proxy program with communication function proposed in this invention differs from the commonly used automated penetration schemes. It further enhances intelligence on the basis of automation, realizing fully intelligent penetration throughout the entire process. At the same time, in response to the problem of the explosive complexity of the action space of reinforcement learning agents in complex large-scale network applications, this invention provides the possibility of cluster operation of agents by applying intelligent encrypted communication.

[0115] One specific embodiment of the present invention is a penetration testing system for attack agents based on reinforcement learning. The system includes a perception module, a decision-making module, a control module, and a communication module, such as... Figure 6 The diagram shows the working framework of the penetration testing process for this system.

[0116] The perception module acts as the eyes of the attacking agent, responsible for acquiring state information about the environment. It uses scanners to perceive the current network composition, monitoring components such as host PCs, servers, routers, and security devices. These sensors capture crucial network data, convert it into feature vectors, and then aggregate the information before passing it to the decision module for further processing.

[0117] The decision-making module acts as the brain of the attacking agent, comprising a decision model and a prior knowledge framework for planning attack strategies. The prior knowledge framework subdivides and generalizes penetration steps and scenario-specific behavioral logic, helping the decision model build basic penetration decision logic for rapid model initialization and providing a stable basic framework for subsequent decision logic. In the decision model, this invention uses a DQN network to process and learn network environment information from the perception module. The model receives feature vectors provided by the perception system, analyzes and processes this data through a deep neural network, finds hidden representations, and applies an attention mechanism to improve computational efficiency. Through learning and optimization, the model predicts the optimal operation to achieve the attack target. In the context of a continuous action space, the decision model prioritizes participant critique algorithms, dynamically adapting to different network conditions and security strategies through agent-environment interaction, thereby improving the success rate and stealth of the attack.

[0118] The control module acts as the hands and feet of the attack agent, responsible for executing the sequence of operations defined by the decision-making module to achieve the attack strategy. It integrates penetration testing tools and a payload generator to ensure the execution of the attack and interaction with the target system. The control module can adapt to different network environments and attack scenarios to improve the success rate and stealth of the attack. When the control system receives an attack action instruction from the decision-making system, it first activates the payload generator to generate the corresponding payload, i.e., a command that directly interacts with the target network. This command, based on the stage of interaction with the target network and the type of vulnerability, invokes the appropriate penetration testing tools to launch an attack on the target system.

[0119] The communication module acts as the mouthpiece for the attacking agents, enabling communication between agents in different scenarios and paving the way for potential cluster attacks. The communication module consists of three parts: a message encoder, a message decoder, and a message aggregator. It can vectorize and encrypt the attacking agents' current strategies and overall strategy style tendencies. By sharing noisy strategy tendency information, attacking agents in different environments can ultimately achieve strategy style sharing and encrypted communication in large-scale network scenarios.

[0120] The reinforcement learning-based attack agent penetration testing system of this invention comprises four key components: perception, decision-making, control, and communication. Through collaborative work, it achieves intelligent and optimized attack strategies. This system is applicable to various network penetration testing scenarios and can be customized and expanded as needed to meet different penetration testing requirements.

[0121] As mentioned earlier, the workflow of a reinforcement learning-based attack agent penetration testing process is as follows: Figure 7 As shown:

[0122] Step 1: The perception module acquires network environment information and then transmits the network environment information to the decision-making module;

[0123] Step 2: The decision-making module combines the prior knowledge framework to determine the current tactical stage of the system and the optimal tactical action for the current stage, and then transmits the decision information to the control module for implementation.

[0124] Step 3: The control module performs corresponding actions based on the decision information from the decision module, and after the action, it restarts the perception module to feed back the new network environment information after the action to the decision module and make a new decision.

[0125] Step 4: After the system has interacted with the network environment a certain number of times, it activates the communication module to communicate with other intelligent agents and encrypts and transmits the network environment features learned by the current intelligent agent to other intelligent agents in the same network.

[0126] Step 5: Exit the work cycle when control of the target system is acquired or the stopping condition is met.

[0127] The foregoing described the reinforcement learning-based penetration testing method and system for attack agents according to embodiments of the present invention. It can be recognized that, compared with traditional AI-based intelligent penetration schemes, embodiments of the present invention improve the intelligence level of agents by introducing reinforcement learning models to achieve more autonomous decision-making capabilities; enable agents to perceive and learn network environment information through information aggregation technology; combine the DQN algorithm and Actor-Critic mechanism with the characteristics of different network environments to make intelligent decisions; simplify dataset dimensions and computation through attention mechanisms and improve the model's sensitivity to key information; quickly initialize the model's behavioral framework through a priori knowledge system to effectively shorten training time and improve environmental adaptability; achieve encrypted communication between agents through an encrypted communication module; and improve training efficiency through the A3C training system, providing possibilities for multi-agent collaborative schemes in complex large-scale networks.

[0128] Reference Figure 8 This invention also provides a computer device, comprising:

[0129] At least one processor;

[0130] At least one memory for storing at least one program;

[0131] When at least one program is executed by at least one processor, such that at least one processor achieves Figure 1 The method shown is a penetration testing method for attack agents based on reinforcement learning.

[0132] The content of the above embodiments of the penetration testing method for attack agents based on reinforcement learning is applicable to the embodiments of this computer device. The specific functions implemented by the embodiments of this computer device are the same as those of the above embodiments of the penetration testing method for attack agents based on reinforcement learning, and the beneficial effects achieved are also the same as those achieved by the above embodiments of the penetration testing method for attack agents based on reinforcement learning.

[0133] This invention also provides a computer-readable storage medium storing a processor-executable program, which, when executed by a processor, is used to implement... Figure 1 The method shown is a penetration testing method for attack agents based on reinforcement learning.

[0134] This invention provides a computer-readable storage medium that can execute the reinforcement learning-based penetration testing method for attack agents provided in the method embodiments of this invention. It can execute any combination of the implementation steps of the method embodiments and has the corresponding functions and beneficial effects of the method.

[0135] This invention also discloses a computer program product or computer program, which includes computer instructions stored in a computer-readable storage medium. A processor of a computer device can read the computer instructions from the computer-readable storage medium, and the processor executes the computer instructions, causing the computer device to perform... Figure 1 The method shown.

[0136] In some alternative embodiments, the functions / operations mentioned in the block diagrams may not occur in the order shown in the operation diagrams. For example, depending on the functions / operations involved, two consecutively shown blocks may actually be executed substantially simultaneously, or the aforementioned blocks may sometimes be executed in reverse order. Furthermore, the embodiments presented and described in the flowcharts of this invention are provided by way of example to provide a more comprehensive understanding of the technology. The disclosed methods are not limited to the operations and logic flows presented herein. Alternative embodiments are contemplated in which the order of various operations is changed and sub-operations described as part of a larger operation are executed independently.

[0137] Furthermore, although the invention has been described in the context of functional modules, it should be understood that, unless otherwise stated, one or more of the aforementioned functions and / or features may be integrated into a single physical device and / or software module, or one or more functions and / or features may be implemented in a separate physical device or software module. It is also understood that a detailed discussion of the actual implementation of each module is unnecessary for understanding the invention. Rather, given the properties, functions, and internal relationships of the various functional modules in the apparatus disclosed herein, the actual implementation of the module will be understood within the scope of conventional skill of an engineer. Therefore, those skilled in the art can implement the invention as set forth in the claims using ordinary techniques without excessive experimentation. It is also understood that the specific concepts disclosed are merely illustrative and not intended to limit the scope of the invention, which is determined by the full scope of the appended claims and their equivalents.

[0138] If the aforementioned functions are implemented as software functional units and sold or used as independent products, they can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of this invention, or the part that contributes to the prior art, or a part of the technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods described in the various embodiments of this invention. The aforementioned storage medium includes various media capable of storing program code, such as USB flash drives, portable hard drives, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0139] The logic and / or steps represented in the flowchart or otherwise described herein, for example, can be considered as a sequenced list of executable instructions for implementing logical functions, and can be embodied in any computer-readable medium for use by, or in conjunction with, an instruction execution system, apparatus, or device (such as a computer-based system, a processor-included system, or other system that can fetch and execute instructions from, an instruction execution system, apparatus, or device). For the purposes of this specification, "computer-readable medium" can be any means that can contain, store, communicate, propagate, or transmit programs for use by, or in conjunction with, an instruction execution system, apparatus, or device.

[0140] More specific examples of computer-readable media (a non-exhaustive list) include: electrical connections (electronic devices) having one or more wires, portable computer disk drives (magnetic devices), random access memory (RAM), read-only memory (ROM), erasable and editable read-only memory (EPROM or flash memory), fiber optic devices, and portable optical disc read-only memory (CDROM). Furthermore, computer-readable media can even be paper or other suitable media on which the aforementioned program can be printed, because the aforementioned program can be obtained electronically, for example, by optically scanning the paper or other medium, followed by editing, interpreting, or otherwise processing as necessary, and then stored in computer memory.

[0141] It should be understood that various parts of the present invention can be implemented in hardware, software, firmware, or a combination thereof. In the above embodiments, multiple steps or methods can be implemented in software or firmware stored in memory and executed by a suitable instruction execution system. For example, if implemented in hardware, as in another embodiment, it can be implemented using any one or a combination of the following techniques known in the art: discrete logic circuits having logic gates for implementing logical functions on data signals, application-specific integrated circuits (ASICs) having suitable combinational logic gates, programmable gate arrays (PGAs), field-programmable gate arrays (FPGAs), etc.

[0142] In the foregoing description of this specification, references to terms such as "one embodiment," "another embodiment," or "some embodiments" indicate that a specific feature, structure, material, or characteristic described in connection with an embodiment or example is included in at least one embodiment or example of the present invention. In this specification, illustrative expressions of the above terms do not necessarily refer to the same embodiment or example. Furthermore, the specific features, structures, materials, or characteristics described may be combined in any suitable manner in one or more embodiments or examples.

[0143] Although embodiments of the invention have been shown and described, those skilled in the art will understand that various changes, modifications, substitutions and alterations can be made to these embodiments without departing from the principles and spirit of the invention, the scope of which is defined by the claims and their equivalents.

[0144] The above is a detailed description of the preferred embodiments of the present invention. However, the present invention is not limited to the embodiments. Those skilled in the art can make various equivalent modifications or substitutions without departing from the spirit of the present invention. All such equivalent modifications or substitutions are included within the scope defined by the claims of this application.

Claims

1. A penetration testing method for attack agents based on reinforcement learning, characterized in that, Includes the following steps: Multiple training thread models are constructed based on the DQN algorithm, the Actor-Critic mechanism, and a pre-defined prior knowledge framework. The multiple training thread models are then trained synchronously using the A3C algorithm to obtain multiple decision models. Obtain network environment information; The current tactical stage is determined based on the network environment information and the decision model, and then the optimal attack action signal is obtained based on the current tactical stage. A simulated attack was launched on the target system based on the optimal attack action signal, and the test results were recorded. Based on the test results, the first current hidden state and message vector of the corresponding decision model are obtained through the message encoder. Then, based on the first current hidden state and the message vector, the policy tendency information is obtained through the message decoder and message aggregator. The other decision models are updated based on the policy tendency information. The prior knowledge framework includes: The penetration testing process is divided into multiple tactical phases, and multiple attack actions are defined according to each tactical phase to obtain a prior knowledge framework. The step of constructing multiple training thread models based on the DQN algorithm, the Actor-Critic mechanism, and a pre-defined prior knowledge framework specifically includes: The reinforcement learning network that applies the attention mechanism is used as the value network Critic, and the pre-defined deep neural network is used as the policy network Actor. The attack action is evaluated using the value network Critic based on the prior knowledge framework to obtain the attack action reward value; The policy network Actor is updated based on the attack action reward value to obtain multiple training thread models; The step of synchronously training multiple training thread models according to the A3C algorithm to obtain multiple decision models specifically includes: Simultaneously, multiple threads are started, and multiple network environment sample information is obtained; Decisions are made based on the network environment sample information described above, and corresponding reward signals are obtained; Update the gradient parameters of the corresponding training thread model according to each of the reward signals; The common model is updated based on the gradient parameters described above, and the number of updates is recorded. When the number of updates exceeds the preset number of updates, each of the training thread models is trained according to the common model to obtain multiple decision models.

2. The penetration testing method for attack agents based on reinforcement learning according to claim 1, characterized in that, The step of launching a simulated attack on the target system based on the optimal attack action signal and recording the test results specifically includes: Based on the optimal attack action signal, a corresponding attack payload is generated by the attack payload generator. Based on the attack payload, a simulated attack is launched on the target system using a penetration testing tool, and the test results are recorded.

3. The penetration testing method for attack agents based on reinforcement learning according to claim 1, characterized in that, The step of obtaining the first current hidden state and message vector of the decision model based on the test results through the message encoder specifically includes: When the test result reaches the preset result, the hidden state of the corresponding decision model at the previous moment is obtained; The previous hidden state is input into the message encoder, and the first current hidden state and the message vector are output.

4. The penetration testing method for an attack agent based on reinforcement learning according to claim 1, characterized in that, The step of obtaining policy bias information based on the first current hidden state and the message vector through a message decoder and a message aggregator, and updating other decision models based on the policy bias information, specifically includes: The first current hidden state and the message vector are input into the message decoder, and the second current hidden state and attention weights are output. The second current hidden state and the attention weight are fused to obtain the policy tendency information; The policy preference information is input into other decision models to update them.

5. A penetration testing system for attack agents based on reinforcement learning, characterized in that: include: The training module is used to construct multiple training thread models based on the DQN algorithm, the Actor-Critic mechanism, and a preset prior knowledge framework, and to synchronously train the multiple training thread models according to the A3C algorithm to obtain multiple decision models. The sensing module is used to acquire network environment information; The decision module is used to determine the current tactical stage based on the network environment information and the decision model, and then obtain the optimal attack action signal based on the current tactical stage. The control module is used to launch a simulated attack on the target system according to the optimal attack action signal and record the test results; The communication module is used to obtain the first current hidden state and message vector of the corresponding decision model through the message encoder based on the test results, and then obtain policy tendency information through the message decoder and message aggregator based on the first current hidden state and the message vector, and update other decision models based on the policy tendency information. The prior knowledge framework includes: The penetration testing process is divided into multiple tactical phases, and multiple attack actions are defined according to each tactical phase to obtain a prior knowledge framework. The step of constructing multiple training thread models based on the DQN algorithm, the Actor-Critic mechanism, and a pre-defined prior knowledge framework specifically includes: The reinforcement learning network that applies the attention mechanism is used as the value network Critic, and the pre-defined deep neural network is used as the policy network Actor. The attack action is evaluated using the value network Critic based on the prior knowledge framework to obtain the attack action reward value; The policy network Actor is updated based on the attack action reward value to obtain multiple training thread models; The step of synchronously training multiple training thread models according to the A3C algorithm to obtain multiple decision models specifically includes: Simultaneously, multiple threads are started, and multiple network environment sample information is obtained; Decisions are made based on the network environment sample information described above, and corresponding reward signals are obtained; Update the gradient parameters of the corresponding training thread model according to each of the reward signals; The common model is updated based on the gradient parameters described above, and the number of updates is recorded. When the number of updates exceeds the preset number of updates, each of the training thread models is trained according to the common model to obtain multiple decision models.

6. A computer device, characterized in that, include: At least one processor; At least one memory for storing at least one program; When the at least one program is executed by the at least one processor, the at least one processor implements the method as described in any one of claims 1-4.

7. A computer-readable storage medium storing a processor-executable program, characterized in that, The processor-executable program, when executed by the processor, is used to implement the method of any one of claims 1-4.

Images