An automated penetration method and system based on Rainbow algorithm

Abstract

The present invention discloses an automated penetration method and system based on the Rainbow algorithm. The method comprises the following steps: obtaining information about the environment to be tested, determining the target network corresponding to the penetration test task and the host information of the target network; determining the attack target and executing the corresponding penetration decision action based on the information about the environment to be tested; constructing a first reward function for the attack target using a first neural network model, and constructing a second reward function for the attack tactics using a second neural network model; calculating a first reward value for the attack target based on the first reward function, and calculating a second reward value for the attack tactics based on the second reward function; dynamically updating the selection strategy for the attack target, and dynamically updating the selection strategy for the attack tactics. The training cost of the present invention is low and the efficiency is high, and the decision accuracy of the model can be improved. The system can be widely applied in the field of computer technology.

Description

Technical Field

[0001] The present invention relates to the field of computer technology, and in particular to an automated penetration method and system based on a Rainbow algorithm. Background Art

[0002] The internet has become an integral part of daily life and work for the vast majority of people worldwide, achieving deep integration and widespread adoption in areas such as information acquisition, social media, and e-commerce. However, as its popularity grows, so too do the security challenges it faces. These include data breaches, online fraud, malware attacks, and privacy issues, all of which pose unprecedented threats to individual users and organizations. From finance and healthcare to education and the military, the ubiquity of the internet has enabled cyber threats to penetrate every aspect of life.

[0003] In today's increasingly complex network environment, penetration testing has become a critical tool for identifying and remediating security vulnerabilities. By simulating hacker attacks, it assesses the security of computer systems, networks, or applications, helping organizations identify potentially overlooked weaknesses so they can be remediated before a real attack occurs, strengthening their ability to withstand real-world attacks. Furthermore, penetration testing can assess the effectiveness of an organization's security policies and response mechanisms, ensuring a swift and effective response to real threats. Penetration testing has become a crucial tool for maintaining network security and protecting organizational assets and information from threats.

[0004] While penetration testing is a crucial component of cybersecurity, it still faces some shortcomings. First, cost is a significant issue. Because penetration testing requires specialized tools and technical support and can impact business operations, it is often costly, placing a heavy burden on resource-constrained businesses and organizations. Second, human resource utilization is a significant challenge. Penetration testing typically requires experienced security experts, but these skills are often in short supply, making it difficult for businesses to find suitable candidates. Reliance on expert expertise is also a potential problem with penetration testing. While experienced experts can provide high-quality test results, different experts may employ varying testing methods and strategies, leading to increased subjectivity and uncertainty. Furthermore, in large, complex networks, manual penetration testing may only focus on certain parts of the network or specific systems, making comprehensive coverage difficult. This can result in critical vulnerabilities going undetected, leaving security risks unresolved, or require more specialized personnel and longer testing times to complete the task. Summary of the Invention

[0005] In view of this, an embodiment of the present invention provides an automated penetration method and system based on the Rainbow algorithm with low training cost and high efficiency to improve decision-making accuracy.

[0006] An aspect of an embodiment of the present invention provides an automated penetration method based on the Rainbow algorithm, comprising the following steps:

[0007] Obtain information about the environment to be tested, determine the target network corresponding to the penetration test task and the host information of the target network;

[0008] Determine the attack target and execute the corresponding penetration decision action based on the test environment information; wherein the penetration decision action includes attack techniques and tactics;

[0009] A first reward function is constructed for the attack target using a first neural network model, and a second reward function is constructed for the attack technique and tactics using a second neural network model;

[0010] Calculating a first reward value for the attack target according to the first reward function, and calculating a second reward value for the attack technique and tactics according to the second reward function;

[0011] The attack target selection strategy is dynamically updated according to the first reward value, and the attack technique and tactic selection strategy is dynamically updated according to the second reward value until the automated penetration process is completed.

[0012] Optionally, the environment information to be tested includes configuration information and vulnerability information of all hosts in the target network;

[0013] The first reward function is constructed for the attack target using the first neural network model, specifically comprising: forming a first feature vector from the environmental state in the measured environmental information and the selected target host, and using the first feature vector as first input information of the first neural network model; inputting the first input information into the first neural network model for training, and completing the construction of the first reward function;

[0014] The second reward function is constructed for the attack tactics using the second neural network model, specifically comprising: forming a second feature vector from the state information of the selected target host, using the second feature vector as second input information of the second neural network model, inputting the second input information into the second neural network model for training, and completing the construction of the second reward function; wherein the state information includes configuration information, open services, operating system, and vulnerability information;

[0015] The architecture of the first reward function and the second reward function is consistent. Each reward function adopts a 4-layer neural network architecture, including an output layer, three hidden layers, and an output layer activated by a softmax function.

[0016] Optionally, the structures of the first neural network model and the second neural network model are consistent;

[0017] The dynamically updating the attack target selection strategy according to the first reward value, and the dynamically updating the attack technique and tactic selection strategy according to the second reward value, comprises the following steps:

[0018] The strategy is updated based on action value, and the off-policy method is used to maintain the Q function of the two neural network models. The learning experience of the two neural network models is used to alternately update the network weights of the other model.

[0019] A distributed RL algorithm is used to construct the Q function, where the distributed RL algorithm replaces the Q value expectation with an estimated Q value distribution;

[0020] Using a multi-step learning mechanism to increase the information accessed by the neural network model;

[0021] Introducing a priority experience replay mechanism to increase the probability of using large error samples for learning updates, thereby increasing the neural network model's attention to decisions with large errors, thereby accelerating the convergence of the neural network model;

[0022] The Dueling Network algorithm is introduced, and two neural network models are used to replace the value function and advantage function in the action value calculation function to improve the stability of Q value estimation;

[0023] Add a noise network mechanism to the neural network model and replace all linear layers with noise networks to enhance the robustness of the model and improve the exploration efficiency of the model;

[0024] The updating of the first neural network model and the second neural network model is completed.

[0025] Optionally, constructing the Q function using a distributed RL algorithm comprises the following steps:

[0026] Split each Q value distribution of the Q function into multiple atoms, and determine the length vector of each atom based on the maximum Q value and the minimum Q value;

[0027] According to the probability mass of the length vector, define the Q value distribution at any time t;

[0028] The input layer of the structure of the Q function is determined as the feature vector in the current state, the middle layer of the Q function is determined as the hidden layer, and the output layer of the Q function is determined as the distribution corresponding to each action value;

[0029] Based on the distribution of each action value, the update objective is optimized to minimize the Kullbeck-Leibler divergence between this distribution and the target distribution.

[0030] Optionally, the multi-step learning mechanism is used to increase the information accessed by the neural network model, comprising the following steps:

[0031] Calculate the learning reward obtained by n steps of learning;

[0032] reconstructing a target distribution and a loss function based on the learned reward;

[0033] The calculation formula for the learning reward is:

[0034]

[0035] in, Represents the cumulative reward obtained by the model agent at time t by looking forward n steps; Represents the learning rate corresponding to the k-th step; k represents the k-th step; n represents the number of steps forward at time t.

[0036] Optionally, the priority experience replay mechanism is introduced to increase the probability of using large error samples for learning updates, thereby increasing the attention of the neural network model to decisions with large target errors, so as to accelerate the convergence of the neural network model. The expression of the sampling probability of the experience sample is:

[0037]

[0038] Among them, p t represents the sampling probability of the sample at time t; ∝ represents the proportional symbol; D KL represents the Kullbeck-Leibler divergence; φ z Representative target distribution Projection onto the first dimension of the distribution; represents the probability distribution of Q values ​​for executing n steps of actions according to the strategy; d t is the target distribution; ω is a hyperparameter that determines the shape of the distribution.

[0039] Optionally, the Dueling Network algorithm is introduced, and two neural network models are used to replace the value function and advantage function parts in the action value calculation function, including the following steps:

[0040] The feature vector formed by the state of the target network is used as the input layer, passes through a shared hidden layer in the middle, and the output of the hidden layer is split and input into the value flow network and the advantage flow network;

[0041] The output results of the value flow network and the advantage flow network are merged, and then normalized using the softmax layer to obtain the final output result.

[0042] Optionally, a noise network mechanism is added to the neural network model, and all linear layers are replaced with noise networks. The expression of the noise network is:

[0043] y=(Wx+b)+((W noisy ⊙∈ ω )x+b noisy ⊙∈ b )

[0044] Where x represents the input vector of the noise network, y represents the output vector of the noise network; W represents the weight parameter to be trained in the noise network; b represents the bias parameter to be trained; W noisy represents the noise weight; b noisy represents noise bias; ⊙ represents element-wise product; ∈ ω represents a random variable; ∈ b represents a random variable.

[0045] Another aspect of the present invention provides an automated infiltration system based on the Rainbow algorithm, including:

[0046] The first module is used to obtain information about the environment to be tested, determine the target network corresponding to the penetration test task and the host information of the target network;

[0047] The second module is used to determine the attack target and execute the corresponding penetration decision action based on the measured environment information; wherein the penetration decision action includes attack techniques and tactics;

[0048] A third module is configured to construct a first reward function for the attack target using the first neural network model, and to construct a second reward function for the attack techniques and tactics using the second neural network model;

[0049] a fourth module, configured to calculate a first reward value for the attack target according to the first reward function, and to calculate a second reward value for the attack technique and tactics according to the second reward function;

[0050] The fifth module is used to dynamically update the selection strategy of the attack target according to the first reward value, and dynamically update the selection strategy of the attack technique and tactics according to the second reward value until the automated penetration process is completed.

[0051] Another aspect of an embodiment of the present invention further provides an electronic device, including a processor and a memory;

[0052] The memory is used to store programs;

[0053] The processor executes the program to implement the method described above.

[0054] The present 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 execute the computer instructions, causing the computer device to perform the above method.

[0055] The embodiment of the present invention obtains information about the environment to be tested, determines the target network corresponding to the penetration test task and the host information of the target network; determines the attack target based on the information about the environment to be tested and executes the corresponding penetration decision action; wherein the penetration decision action includes attack techniques and tactics; uses a first neural network model to construct a first reward function for the attack target, and uses a second neural network model to construct a second reward function for the attack techniques and tactics; calculates a first reward value for the attack target based on the first reward function, and calculates a second reward value for the attack techniques and tactics based on the second reward function; dynamically updates the selection strategy for the attack target based on the first reward value, and dynamically updates the selection strategy for the attack techniques and tactics based on the second reward value, until the automated penetration process is completed. The training cost of the present invention is low and the efficiency is high, and it can improve the decision accuracy of the model. BRIEF DESCRIPTION OF THE DRAWINGS

[0056] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the following briefly introduces the drawings required for use in the description of the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work.

[0057] Figure 1 An overall step flow chart provided for an embodiment of the present invention;

[0058] Figure 2 A schematic diagram of the overall framework provided by an embodiment of the present invention;

[0059] Figure 3 Schematic diagram of the distributed RL combined with Dueling network architecture provided by an embodiment of the present invention;

[0060] Figure 4A schematic structural diagram of an electronic device provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0061] In order to make the objects, technical solutions and advantages of the present invention more clearly understood, the present invention is further described in detail below with reference to the accompanying drawings and examples. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention. When the following description refers to the drawings, unless otherwise indicated, the same numbers in different drawings represent the same or similar elements. The embodiments described in the following exemplary embodiments do not represent all embodiments consistent with the embodiments of the present invention. They are merely examples of devices and methods consistent with some aspects of the embodiments of the present invention as detailed in the appended claims.

[0062] It will be understood that the terms "first," "second," and the like used in the present invention may be used herein to describe various concepts, but unless otherwise specified, these concepts are not limited by these terms. These terms are merely used to distinguish one concept from another. For example, without departing from the scope of the embodiments of the present invention, first information may also be referred to as second information, and similarly, second information may also be referred to as first information. Depending on the context, the words "if" and "if" as used herein may be interpreted as "at the time of," "when," or "in response to a determination."

[0063] The terms "at least one", "plurality", "each", "any", etc. used in the present invention include at least one, two or more, multiple, two or more, each refers to each of the corresponding multiple, and any refers to any one of the multiple.

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

[0065] The automated penetration method based on the Rainbow algorithm provided in an embodiment of the present invention relates to the field of computer technology. The automated penetration method based on the Rainbow algorithm provided in an embodiment of the present invention can be applied to a terminal, can also be applied to a server, and can also be software running in a terminal or a server. In some embodiments, the terminal can be a smart phone, a tablet computer, a laptop computer, a desktop computer, a smart speaker, a smart watch, and a car terminal, etc., but is not limited to this; the server side can be configured as an independent physical server, or as a server cluster or distributed system composed of multiple physical servers, or as a cloud server that provides basic cloud computing services such as cloud services, cloud databases, cloud computing, cloud functions, cloud storage, network services, cloud communications, middleware services, domain name services, security services, CDN, and big data and artificial intelligence platforms. The server can also be a node server in a blockchain network; the software can be an application that implements the automated penetration method based on the Rainbow algorithm, etc., but is not limited to the above forms.

[0066] The present invention can be used in a wide variety of general or special computer system environments or configurations. For example: personal computers, server computers, handheld or portable devices, tablet devices, multiprocessor systems, microprocessor-based systems, set-top boxes, programmable consumer electronics, network PCs, minicomputers, mainframe computers, distributed computing environments including any of the above systems or devices, and the like. The present invention can be described in the general context of computer-executable instructions executed by a computer, such as program modules. Generally, program modules include routines, programs, objects, components, data structures, etc. that perform specific tasks or implement specific abstract data types. The present invention can also be practiced in distributed computing environments in which tasks are performed by remote processing devices connected via a communications network. In a distributed computing environment, program modules can be located in local and remote computer storage media, including storage devices.

[0067] Automated penetration testing has become a new hotspot in related technologies due to the significant human, material, and financial constraints of manual penetration testing. Compared to time-consuming and costly manual penetration testing, automated testing can cover a wider range of network areas in a shorter timeframe, significantly improving the efficiency and speed of penetration testing, reducing the overall cost of penetration testing and making continuous security assessments economically viable. Furthermore, the ease of use and rapid update of automated tools significantly reduce the reliance on expert experience and the technical threshold required for penetration testing. They can also identify and test the latest security threats, keeping security measures current and effective. Furthermore, automated penetration testing provides a consistent testing methodology that can be repeated across different environments to ensure the effectiveness and continuity of security measures and reduce human error.

[0068] In recent years, a popular approach has been using reinforcement learning (RL) to solve the action decision-making and path selection problems in automated penetration testing. Trained agents can make action decisions based on environmental conditions, significantly improving the efficiency of penetration testing. For example, researchers have characterized the penetration testing process as a Markov decision process (MDP), defining the transition probability from each moment to the current state as dependent solely on the current state and action, independent of previous states and actions. This framework establishes a state and action description framework for the penetration testing process and employs the Q-MDP algorithm to solve the model. Since then, many researchers have improved upon this framework. The Q-Learning algorithm is widely used to update policies. Q-learning updates policies based on the difference between the optimal updated estimate of the Q value and the current estimate. Compared to indirectly solving for the optimal policy through a value function, the policy gradient method directly differentiates the objective reward function to obtain the optimal gradient direction. Gradient descent is then used to obtain the optimal policy parameters, thereby achieving the optimal policy, accelerating convergence. On this basis, some technicians have introduced the advantage function to improve the stability of training, and modeled the penetration testing task as a partially observable Markov decision process (POMDP) ​​problem. The model agent can perform testing without knowing the situation of the entire network, and use an external POMDP solver to solve the policy update. Although uncertainty is taken into account, it is still not feasible for existing solvers to solve POMDP in a large network.

[0069] With the development of deep reinforcement learning (DRL) technology, new breakthroughs have been made in the field of automated penetration testing. Many researchers have combined the DQN algorithm to build models. The open-source project DeepExploit integrates automated penetration tools such as Nmap and Metasploit, using the A3C algorithm to achieve optimal path search. Other researchers have proposed a new algorithm, called NIG-AP, which incorporates network information gain into attack planning and proposes an attack planning algorithm based on network information gain to plan attack paths within the target network. Others have employed hierarchical reinforcement learning algorithms, breaking down the RL problem into several subproblems to address the issue of excessively large state spaces. Some researchers have improved the reward function by incorporating human feedback into the reward model to improve decision-making accuracy. Alternatively, they have adopted the concept of imitation learning, incorporating expert prior knowledge into the reward function.

[0070] In the methods discussed above, the agent's behavior is designed based on the combination of target hosts and attack strategies (including scanning, vulnerability exploitation, and privilege escalation). Because each target host may employ a different attack combination, the number of executable actions is the number of target hosts multiplied by the number of attack combinations. Therefore, when the network environment is complex, the number of actions that the agent can perform becomes very large. This design means that the model agent may need to make a large number of attempts before obtaining effective feedback. Furthermore, in large-scale networks, due to the large number of devices and complex network topology, the state space can also be excessively large. These factors can lead to low sample efficiency in reinforcement learning algorithms, requiring more training data and experiments to achieve good performance. The exploration efficiency and intelligence of existing models need to be improved.

[0071] Currently, automated penetration testing methods based on traditional RL algorithms typically require large amounts of training data to achieve good performance. This is especially true when faced with complex penetration testing scenarios, which can require numerous simulations, resulting in lengthy model training times and significant computational resources. Furthermore, while automated penetration testing methods implemented using RL algorithms in conjunction with deep neural networks exhibit superior intelligence and more efficient exploration, they also suffer from slow convergence, high computing power requirements, and high training costs when faced with large networks. Furthermore, neural networks are prone to overestimating action and state values ​​during policy updates, leading to large fluctuations in estimates and network instability, which are issues that need to be addressed in future improvements.

[0072] Based on the above problems, the present invention provides an automated penetration method based on the Rainbow algorithm, aiming to solve the following problems:

[0073] ① To address the problem of low exploration efficiency caused by excessively large state space and action space, this paper proposes a new model representation framework, which further decomposes and subdivides the penetration process, decouples the action space, separates the target host selection and attack method selection in the action design of traditional schemes, and replaces the single selection model with two selection models. This makes the actions and decisions of the model agent closer to the expert decision-making in real scenarios, thereby improving the intelligence of the model.

[0074] ② Currently, existing automated penetration methods based on reinforcement learning have problems such as slow convergence and low learning efficiency of experience samples. This paper introduces the Rainbow algorithm to train the model. The Rainbow algorithm integrates multiple improved reinforcement learning algorithms, among which the priority experience replay algorithm enables the intelligent agent to prioritize learning experiences with high expected learning value, thereby improving learning efficiency. The combination of Double Q-Learning, distributed RL algorithm and Dueling network algorithm achieves more accurate and stable value estimation. The addition of a noise network algorithm introduces an exploration mechanism by adding noise to the network parameters, thereby improving exploration efficiency. The multi-step learning mechanism enables the model to update with a larger step, further accelerating the learning speed.

[0075] ③At the same time, the reward function of existing penetration methods based on reinforcement learning algorithms is usually designed by experts through trial and error. For the same algorithm, choosing different reward functions will result in different performance differences. Therefore, the design of reward functions has always been a difficult problem for RL algorithms. The present invention uses a neural network to fit the reward function and train it with expert experience as a label, instead of manually designing the reward function. This realizes the self-construction of the reward function and uses the neural network to approximate the optimal reward function. It has higher flexibility, adaptability and generalization ability, can better adapt to various environments and tasks, and can handle some complex tasks that are difficult to handle with traditional methods.

[0076] Specifically, if Figure 1 As shown, an embodiment of the present invention provides an automated penetration method based on the Rainbow algorithm, comprising the following steps:

[0077] Obtain information about the environment to be tested, determine the target network corresponding to the penetration test task and the host information of the target network;

[0078] Determine the attack target and execute the corresponding penetration decision action based on the test environment information; wherein the penetration decision action includes attack techniques and tactics;

[0079] A first reward function is constructed for the attack target using a first neural network model, and a second reward function is constructed for the attack technique and tactics using a second neural network model;

[0080] Calculating a first reward value for the attack target according to the first reward function, and calculating a second reward value for the attack technique and tactics according to the second reward function;

[0081] The attack target selection strategy is dynamically updated according to the first reward value, and the attack technique and tactic selection strategy is dynamically updated according to the second reward value until the automated penetration process is completed.

[0082] Optionally, the environment information to be tested includes configuration information and vulnerability information of all hosts in the target network;

[0083] The first reward function is constructed for the attack target using the first neural network model, specifically comprising: forming a first feature vector from the environmental state in the measured environmental information and the selected target host, and using the first feature vector as first input information of the first neural network model; inputting the first input information into the first neural network model for training, and completing the construction of the first reward function;

[0084] The second reward function is constructed for the attack tactics using the second neural network model, specifically comprising: forming a second feature vector from the state information of the selected target host, using the second feature vector as second input information of the second neural network model, inputting the second input information into the second neural network model for training, and completing the construction of the second reward function; wherein the state information includes configuration information, open services, operating system, and vulnerability information;

[0085] The architecture of the first reward function and the second reward function is consistent. Each reward function adopts a 4-layer neural network architecture, including an output layer, three hidden layers, and an output layer activated by a softmax function.

[0086] Optionally, the structures of the first neural network model and the second neural network model are consistent;

[0087] The dynamically updating the attack target selection strategy according to the first reward value, and the dynamically updating the attack technique and tactic selection strategy according to the second reward value, comprises the following steps:

[0088] The strategy is updated based on action value, and the off-policy method is used to maintain the Q function of the two neural network models. The learning experience of the two neural network models is used to alternately update the network weights of the other model.

[0089] A distributed RL algorithm is used to construct the Q function, where the distributed RL algorithm replaces the Q value expectation with an estimated Q value distribution;

[0090] Using a multi-step learning mechanism to increase the information accessed by the neural network model;

[0091] Introducing a priority experience replay mechanism to increase the probability of using large error samples for learning updates, thereby increasing the neural network model's attention to decisions with large errors, thereby accelerating the convergence of the neural network model;

[0092] The Dueling Network algorithm is introduced, and two neural network models are used to replace the value function and advantage function in the action value calculation function to improve the stability of Q value estimation;

[0093] Add a noise network mechanism to the neural network model and replace all linear layers with noise networks to enhance the robustness of the model and improve the exploration efficiency of the model;

[0094] The updating of the first neural network model and the second neural network model is completed.

[0095] Optionally, constructing the Q function using a distributed RL algorithm comprises the following steps:

[0096] Split each Q value distribution of the Q function into multiple atoms, and determine the length vector of each atom based on the maximum Q value and the minimum Q value;

[0097] According to the probability mass of the length vector, define the Q value distribution at any time t;

[0098] The input layer of the structure of the Q function is determined as the feature vector in the current state, the middle layer of the Q function is determined as the hidden layer, and the output layer of the Q function is determined as the distribution corresponding to each action value;

[0099] Based on the distribution of each action value, the update objective is optimized to minimize the Kullbeck-Leibler divergence between this distribution and the target distribution.

[0100] Optionally, the multi-step learning mechanism is used to increase the information accessed by the neural network model, comprising the following steps:

[0101] Calculate the learning reward obtained by n steps of learning;

[0102] reconstructing a target distribution and a loss function based on the learned reward;

[0103] The calculation formula of the learning reward is:

[0104]

[0105] in, Represents the cumulative reward obtained by the model agent at time t by looking forward n steps; Represents the learning rate corresponding to the k-th step; k represents the k-th step; n represents the number of steps forward at time t.

[0106] Optionally, the priority experience replay mechanism is introduced to increase the probability of using large error samples for learning updates, thereby increasing the attention of the neural network model to decisions with large target errors, so as to accelerate the convergence of the neural network model. The expression of the sampling probability of the experience sample is:

[0107]

[0108] Among them, p t represents the sampling probability of the sample at time t; ∝ represents the proportional symbol; D KL represents the Kullbeck-Leibler divergence; φ z Representative target distribution Projection onto the first dimension of the distribution; represents the probability distribution of Q values ​​for executing n steps of actions according to the strategy; d t is the target distribution; ω is a hyperparameter that determines the shape of the distribution.

[0109] Optionally, the Dueling Network algorithm is introduced, and two neural network models are used to replace the value function and advantage function parts in the action value calculation function, including the following steps:

[0110] The feature vector formed by the target network state to be tested is used as the input layer, passes through a shared hidden layer in the middle, and the output of the hidden layer is split and input into the value flow network and the advantage flow network;

[0111] The output results of the value flow network and the advantage flow network are merged, and then normalized using the softmax layer to obtain the final output result.

[0112] Optionally, a noise network mechanism is added to the neural network model, and all linear layers are replaced with noise networks. The expression of the noise network is:

[0113] y=(Wx+b)+((W noisy ⊙∈ ω )x+b noisy ⊙∈ b )

[0114] Where x represents the input vector of the noise network, y represents the output vector of the noise network; W represents the weight parameter to be trained in the noise network; b represents the bias parameter to be trained; W noisy represents the noise weight; b noisy represents noise bias; ⊙ represents element-wise product; ∈ ω represents a random variable; ∈ b represents a random variable.

[0115] The following describes the specific implementation process of the present invention in detail using a specific application scenario as an example:

[0116] like Figure 2 As shown, the core framework of the present invention consists of four parts: the test environment, decision actions, reward function, and policy update. The test environment is the target network for the penetration test task, and the environment state is defined as the configuration details and vulnerability information of all hosts observed by the agent in the test network environment. The decision action component splits the penetration action into two parts: target host selection and attack technique and tactic selection. The reward function uses two neural network models to construct reward functions for host selection and attack technique and tactic selection, respectively. The reward function for host selection takes as input a feature vector consisting of the current environment state and the selected host, and outputs a reward value. The reward function for attack technique and tactic selection takes as input the specific configuration and vulnerability information of the selected host, and outputs a reward value. State-action pairs used during DeepExploit execution are automatically collected in advance as training samples for the two reward functions to train the model. The policy update component uses the Rainbow algorithm, based on the DQN algorithm framework, and integrates prioritized experience replay, double QLearning, dueling DQN, prioritized experience replay, multi-step learning, noise networks, and distributed Q functions to improve the model's exploration capability, learning efficiency, and convergence speed.

[0117] The following describes in detail the policy update algorithm of the embodiment of the present invention:

[0118] This paper adopts the Rainbow algorithm and integrates multiple reinforcement learning improvement algorithms on the basis of the DQN algorithm to improve learning stability and efficiency, increase exploration efficiency, reduce over-estimation bias, and provide more accurate risk assessment.

[0119] Specifically, in this application scenario, the policy update algorithm includes the following steps:

[0120] S21: The present invention adopts an action-value-based approach to update the Agent strategy, using an off-policy method to maintain two Q functions to avoid overestimation of the Q value. The two Q functions are two neural networks with exactly the same structure, and use each other's experience to alternately update the network weights θ and

[0121] S22: The construction of Q function adopts distributed RL algorithm, which replaces Q value expectation with estimated Q value distribution, which further enhances the versatility of Q function and makes the predicted value of function closer to the real Q value, which can greatly improve the convergence speed of model. atoms atoms, and define the length as N atoms A vector z, where V min ,V max They correspond to the minimum Q value and the maximum Q value respectively. The probability mass of z corresponds to the same length N. atoms The vector p θ ,in Corresponding to the distribution z i The probability of occurrence. From this, we can define the distribution d at time t t :

[0122] d t =(z,p θ (S t ,A t ))

[0123] The structure of the Q function is that the input layer is the feature vector of the current state s, the middle layer is the hidden layer, and the output layer is the distribution corresponding to each action value. The optimization goal is to update θ and minimize the distribution d t and target distribution The Kullbeck-Leibler divergence D between KL (φ z d′t||d t ), so that this distribution closely matches the actual return distribution.

[0124] in, Indicates that in S t+1 The action with the largest action value in the state, φ z is the L2 projection of the target distribution on z.

[0125] S23: To increase learning speed and save computational costs, a multi-step learning mechanism is used, allowing the model to explore more situations in the future and improve learning accuracy. The reward for n-step learning is:

[0126]

[0127] Furthermore, we reconstruct the target distribution and loss function:

[0128]

[0129] S24: During the network update process, the errors of some collected experience samples may differ significantly from the estimated values. Methods based on exploration mechanisms or random sampling for updates require multiple rounds of updates to reduce these errors, resulting in low model learning efficiency. To improve the learning efficiency of the model, a priority experience replay mechanism is introduced to increase the probability of using samples with larger errors for learning updates, increase the model's attention to current decisions with larger errors, and accelerate model convergence. The sampling probability of the experience sample is defined as follows:

[0130]

[0131] Here, ω is a hyperparameter that determines the shape of the distribution. At the same time, newly collected experience samples are given the maximum probability to ensure that the model learns new knowledge.

[0132] S25: In order to make a more stable estimate of the Q value, the present invention introduces the Dueling Network algorithm, whose basic idea is another form of the action value calculation function Replace the value function V with two neural networks * (S) and advantage function A * (s,a) part.

[0133] Architecture such as Figure 3 As shown, the input layer is the feature vector composed of state s, which passes through a shared hidden layer f ξ (s), then f ξ The output of (s) is split, and the value flow v η (f ξ The output length of (s) is N atoms vector, dominant flow a ψ 9f ξ The output of (s) is N atoms ×|A| tensor, where represents the advantage value of the i-th atom in the Q distribution corresponding to taking action a in state s. For each atom z i , merge the value stream and the advantage stream, then use the softmax layer to normalize and finally output.

[0134]

[0135] Where φ = f ξ (s),

[0136] S26: In order to enhance the robustness of the model and improve the exploration efficiency of the model, the present invention adds a noise network mechanism to the neural network and replaces all linear layers y=Wx+b with the following noise network:

[0137] y=(Wx+b)+((W noisy ⊙∈ ω )x+b noisy ⊙∈ b ),

[0138] Among them, ∈ ω and ∈ b are all random variables, and ⊙ represents the product of elements.

[0139] As the training process progresses, the model will gradually ignore the noise and achieve automatic annealing, making the model more stable without affecting the accuracy of the optimal decision.

[0140] In summary, the present invention proposes an automatic penetration method based on the Rainbow algorithm to address the problems of poor adaptability, low learning efficiency, slow convergence speed, and high training cost in existing automatic penetration methods based on reinforcement learning algorithms. First, compared with the existing model framework design, the present invention separates the target host selection and attack technique and tactical combination in the decision-making action, greatly reducing the dimension of the action space and being closer to the manual decision-making steps in real scenarios. Secondly, in order to further improve the degree of automation and the generalization ability of the framework, the present invention uses a neural network to fit the reward function, uses pre-collected expert action samples as labels for training, realizes the automatic construction of the reward function, and reduces human intervention. Finally, the Rainbow algorithm is introduced to train the Agent, and the convergence speed and learning efficiency of the model training are improved through the multi-step learning, priority experience replay, noise network and other improved algorithms integrated in the Rainbow algorithm to reduce the training cost.

[0141] Compared with the prior art, the present invention has the following advantages:

[0142] ① This paper proposes a new automated penetration method framework, which divides the penetration decision-making process into two parts: target machine selection and attack method decision-making, greatly reducing the dimension of the action space, and solving the problems of current methods such as high training cost, difficulty in action value calculation, slow model convergence speed, and lack of model accuracy due to the excessive number of action combinations.

[0143] ② This paper uses the Rainbow algorithm to construct the RL model and introduces multiple data augmentation methods to further improve the model's training speed. It utilizes a double Q-learning approximation to reduce the impact of overestimation bias, employs a prioritized replay experience sampling mechanism with higher learning potential, and uses multi-step learning to increase learning speed. Furthermore, it integrates decision networks, distributed RL, and noisy networks to improve the quality of learning probability, resulting in significant improvements in accuracy and convergence speed.

[0144] ③ The present invention uses DNN fitting reward function instead of manually designed reward function, and uses the state+action-reward pairs of expert prior knowledge as labels to train the reward function network, realize the automatic construction of reward function, and improve the versatility and adaptability of the model.

[0145] The embodiment of the present invention further provides an automated infiltration system based on the Rainbow algorithm, comprising:

[0146] The first module is used to obtain information about the environment to be tested, determine the target network corresponding to the penetration test task and the host information of the target network;

[0147] The second module is used to determine the attack target and execute the corresponding penetration decision action based on the measured environment information; wherein the penetration decision action includes attack techniques and tactics;

[0148] A third module is configured to construct a first reward function for the attack target using the first neural network model, and to construct a second reward function for the attack techniques and tactics using the second neural network model;

[0149] a fourth module, configured to calculate a first reward value for the attack target according to the first reward function, and to calculate a second reward value for the attack technique and tactics according to the second reward function;

[0150] The fifth module is used to dynamically update the selection strategy of the attack target according to the first reward value, and dynamically update the selection strategy of the attack technique and tactics according to the second reward value until the automated penetration process is completed.

[0151] It can be understood that the contents of the above method embodiments are all applicable to the present system embodiments, the functions specifically implemented by the present system embodiments are the same as those of the above method embodiments, and the beneficial effects achieved are also the same as those achieved by the above method embodiments.

[0152] An embodiment of the present invention further provides an electronic device comprising a memory and a processor. The memory stores a computer program, and the processor executes the computer program to implement the automated penetration method based on the Rainbow algorithm. The electronic device can be any intelligent terminal, including a tablet computer and an in-vehicle computer.

[0153] It can be understood that the contents of the above method embodiments are applicable to the present device embodiments, the functions specifically implemented by the present device embodiments are the same as those of the above method embodiments, and the beneficial effects achieved are also the same as those achieved by the above method embodiments.

[0154] See also Figure 4 , Figure 4 The hardware structure of an electronic device according to another embodiment is shown. The electronic device includes:

[0155] The processor 401 may be implemented as a general-purpose CPU (Central Processing Unit), a microprocessor, an application-specific integrated circuit (ASIC), or one or more integrated circuits, and is configured to execute relevant programs to implement the technical solutions provided by the embodiments of the present invention.

[0156] The memory 402 can be implemented in the form of a read-only memory (ROM), a static storage device, a dynamic storage device, or a random access memory (RAM). The memory 402 can store an operating system and other application programs. When the technical solutions provided in the embodiments of this specification are implemented through software or firmware, the relevant program codes are stored in the memory 402 and are called by the processor 401 to execute the automated penetration method based on the Rainbow algorithm of the embodiment of the present invention.

[0157] Input / output interface 403, used to implement information input and output;

[0158] Communication interface 404, used to implement communication interaction between this device and other devices, which can be achieved through wired means (such as USB, network cable, etc.) or wireless means (such as mobile network, WiFi, Bluetooth, etc.);

[0159] Bus 405 , which transmits information between various components of the device (e.g., processor 401 , memory 402 , input / output interface 403 , and communication interface 404 );

[0160] The processor 401 , the memory 402 , the input / output interface 403 and the communication interface 404 are connected to each other in communication within the device via a bus 405 .

[0161] An embodiment of the present invention further provides a computer-readable storage medium storing a computer program. When the computer program is executed by a processor, the computer program implements the above-mentioned automated infiltration method based on the Rainbow algorithm.

[0162] It can be understood that the contents of the above method embodiments are all applicable to the present storage medium embodiment, the functions specifically implemented by the present storage medium embodiment are the same as those of the above method embodiments, and the beneficial effects achieved are also the same as those achieved by the above method embodiments.

[0163] The memory, as a non-transient computer-readable storage medium, can be used to store non-transient software programs and non-transient computer executable programs. In addition, the memory may include a high-speed random access memory and may also include a non-transient memory, such as at least one disk storage device, a flash memory device, or other non-transient solid-state storage device. In some embodiments, the memory may optionally include a memory remotely arranged relative to the processor, and these remote memories may be connected to the processor via a network. Examples of the above-mentioned network include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0164] In some optional embodiments, the function / operation mentioned in the block diagram may not occur in the order mentioned in the operation diagram. For example, depending on the function / operation involved, the two boxes shown in succession can actually be executed substantially simultaneously or the boxes can sometimes be executed in reverse order. In addition, the embodiment presented and described in the flow chart of the present invention is provided in an exemplary manner for the purpose of providing a more comprehensive understanding of the technology. The disclosed method is not limited to the operation and logic flow presented herein. Optional embodiments are contemplated in which the order of the various operations is changed and the sub-operations described as a part of a larger operation are performed independently.

[0165] Furthermore, although the present invention is described in the context of functional modules, it should be understood that, unless otherwise indicated, one or more of the functions and / or features described may be integrated into a single physical device and / or software module, or one or more functions and / or features may be implemented in separate physical devices or software modules. It will also be understood that a detailed discussion of the actual implementation of each module is not necessary for understanding the present invention. More specifically, given the properties, functions, and internal relationships of the various functional modules in the devices disclosed herein, the actual implementation of the module will be understood within the ordinary skill of an engineer. Therefore, a person skilled in the art using ordinary skill will be able to implement the present invention set forth in the claims without undue experimentation. It will also be understood that the specific concepts disclosed are merely illustrative and are not intended to limit the scope of the present invention, which is determined by the full scope of the appended claims and their equivalents.

[0166] If the functions are implemented in the form of 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 the present invention, or the part that contributes to the prior art, or the part of the technical solution, can be embodied in the form of a software product. The computer software product is stored in a storage medium and includes several instructions for enabling a computer device (which can be a personal computer, a server, or a network device, etc.) to execute all or part of the steps of the method described in each embodiment of the present invention. The aforementioned storage medium includes various media that can store program codes, such as a USB flash drive, a mobile hard disk, a read-only memory (ROM), a random access memory (RAM), a magnetic disk, or an optical disk.

[0167] The logic and / or steps represented in the flowcharts or otherwise described herein, for example, can be considered as an ordered list of executable instructions for implementing the 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 (e.g., a computer-based system, a system including a processor, or other system that can fetch and execute instructions from an instruction execution system, apparatus, or device). For purposes of this specification, a "computer-readable medium" can be any device that can contain, store, communicate, propagate, or transport a program for use by, or in conjunction with, an instruction execution system, apparatus, or device.

[0168] More specific examples (a non-exhaustive list) of computer-readable media include the following: an electrical connection with one or more wires (electronic devices), a portable computer disk cartridge (magnetic devices), a random access memory (RAM), a read-only memory (ROM), an erasable and programmable read-only memory (EPROM or flash memory), a fiber optic device, and a portable compact disc read-only memory (CDROM). In addition, the computer-readable medium may even be paper or other suitable medium on which the program is printed, since the program may be obtained electronically, for example, by optically scanning the paper or other medium, followed by editing, deciphering, or processing in another suitable manner as necessary, and then stored in a computer memory.

[0169] It should be understood that various parts of the present invention can be implemented using hardware, software, firmware, or a combination thereof. In the above-described embodiments, multiple steps or methods can be implemented using software or firmware stored in a memory and executed by a suitable instruction execution system. For example, if implemented using hardware, as in another embodiment, any one of the following technologies known in the art or a combination thereof can be used: a discrete logic circuit having a logic gate circuit for implementing a logic function on a data signal, an application-specific integrated circuit having a suitable combination of logic gate circuits, a programmable gate array (PGA), a field programmable gate array (FPGA), etc.

[0170] Throughout this specification, reference to terms such as "one embodiment," "some embodiments," "examples," "specific examples," or "some examples" means that a specific feature, structure, material, or characteristic described in conjunction with that embodiment or example is included in at least one embodiment or example of the present invention. In this specification, schematic representations 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 any one or more embodiments or examples.

[0171] While embodiments of the present invention have been shown and described, it will be appreciated by those skilled in the art that various changes, modifications, substitutions, and variations may be made to the embodiments without departing from the principles and spirit of the invention, and that the scope of the invention is defined by the claims and their equivalents.

[0172] The above is a specific description of the preferred implementation of the present invention, but the present invention is not limited to the embodiments. Those skilled in the art can make various equivalent modifications or substitutions without violating the spirit of the present invention. These equivalent modifications or substitutions are all included in the scope defined by the claims of the present invention.

Claims

1. An automated penetration method based on the Rainbow algorithm, characterized in that: The following steps are involved: Obtain information about the environment to be tested, determine the target network corresponding to the penetration test task and the host information of the target network; Determine the attack target and execute the corresponding penetration decision action based on the test environment information; wherein the penetration decision action includes attack techniques and tactics; A first reward function is constructed for the attack target using a first neural network model, and a second reward function is constructed for the attack technique and tactics using a second neural network model; Calculating a first reward value for the attack target according to the first reward function, and calculating a second reward value for the attack technique and tactics according to the second reward function; Dynamically updating the attack target selection strategy according to the first reward value, and dynamically updating the attack technique and tactic selection strategy according to the second reward value, until the automated penetration process is completed; The structures of the first neural network model and the second neural network model are consistent; The dynamically updating the attack target selection strategy according to the first reward value, and the dynamically updating the attack technique and tactic selection strategy according to the second reward value, comprises the following steps: The strategy is updated based on action value, and the off-policy method is used to maintain the Q function of the two neural network models. The learning experience of the two neural network models is used to alternately update the network weights of the other model. A distributed RL algorithm is used to construct the Q function, where the distributed RL algorithm replaces the Q value expectation with an estimated Q value distribution; Using a multi-step learning mechanism to increase the information accessed by the neural network model; Introducing a priority experience replay mechanism to speed up the convergence of the neural network model; The Dueling Network algorithm is introduced, and two neural network models are used to replace the value function and advantage function in the action value calculation function to improve the stability of Q value estimation; Add a noise network mechanism to the neural network model and replace all linear layers with noise networks to enhance the robustness of the model and improve the exploration efficiency of the model; The updating of the first neural network model and the second neural network model is completed.

2. The automated penetration method based on the Rainbow algorithm according to claim 1, characterized in that: The environment information to be tested includes configuration information and vulnerability information of all hosts in the target network; The first reward function is constructed for the attack target using the first neural network model, specifically comprising: forming a first feature vector from the environmental state in the measured environmental information and the selected target host, and using the first feature vector as first input information of the first neural network model; inputting the first input information into the first neural network model for training, and completing the construction of the first reward function; The second reward function is constructed for the attack tactics using the second neural network model, specifically comprising: forming a second feature vector from the state information of the selected target host, using the second feature vector as second input information of the second neural network model, inputting the second input information into the second neural network model for training, and completing the construction of the second reward function; wherein the state information includes configuration information, open services, operating system, and vulnerability information; The architecture of the first reward function and the second reward function is consistent. Each reward function adopts a 4-layer neural network architecture, including an output layer, three hidden layers, and an output layer activated by a softmax function.

3. The automated penetration method based on the Rainbow algorithm according to claim 1, characterized in that: The Q function is constructed using a distributed RL algorithm, including the following steps: Split each Q value distribution of the Q function into multiple atoms, and determine the length vector of each atom based on the maximum Q value and the minimum Q value; According to the probability mass of the length vector, define the Q value distribution at any time t; The input layer of the structure of the Q function is determined as the feature vector in the current state, the middle layer of the Q function is determined as the hidden layer, and the output layer of the Q function is determined as the distribution corresponding to each action value; Based on the distribution of each action value, the update objective is optimized to minimize the Kullbeck-Leibler divergence between this distribution and the target distribution.

4. The automated penetration method based on the Rainbow algorithm according to claim 1, characterized in that: The multi-step learning mechanism is used to increase the information accessed by the neural network model, including the following steps: Calculate the learning reward obtained by n steps of learning; reconstructing a target distribution and a loss function based on the learned reward; The calculation formula of the learning reward is: in, Represents the cumulative reward obtained by the model agent at time t by looking forward n steps; Represents the learning rate corresponding to the k-th step; k represents the k-th step; n represents the number of steps forward at time t.

5. The automated penetration method based on the Rainbow algorithm according to claim 1, characterized in that: The priority experience replay mechanism is introduced to accelerate the convergence of the neural network model, where the sampling probability of the experience sample is expressed as: Among them, p t represents the sampling probability of the sample at time t; ∝ represents the proportional symbol; D KL represents the Kullbeck-Leibler divergence; φ z Representative target distribution Projection onto the first dimension of the distribution; represents the probability distribution of Q values ​​for executing n steps of actions according to the strategy; d t is the target distribution; ω is a hyperparameter that determines the shape of the distribution.

6. The automated penetration method based on the Rainbow algorithm according to claim 1, characterized in that: The Dueling Network algorithm is introduced, and two neural network models are used to replace the value function and advantage function parts in the action value calculation function, including the following steps: The feature vector formed by the state of the target network is used as the input layer, passes through a shared hidden layer in the middle, and the output of the hidden layer is split and input into the value flow network and the advantage flow network; The output results of the value flow network and the advantage flow network are merged, and then normalized using the softmax layer to obtain the final output result.

7. The automated penetration method based on the Rainbow algorithm according to claim 1, characterized in that: The noise network mechanism is added to the neural network model, and all linear layers are replaced with the noise network. The expression of the noise network is: y=(Wx+b)+((W noisy ⊙∈ ω )x+b noisy ⊙∈ b ) Where x represents the input vector of the noise network, y represents the output vector of the noise network; W represents the weight parameter to be trained in the noise network; b represents the bias parameter to be trained; W noisy represents the noise weight; b noisy represents noise bias; ⊙ represents element-wise product; ∈ ω represents a random variable; ∈ b represents a random variable.

8. A system for implementing the automated infiltration method based on the Rainbow algorithm according to any one of claims 1 to 7, characterized in that: include: The first module is used to obtain information about the environment to be tested, determine the target network corresponding to the penetration test task and the host information of the target network; The second module is used to determine the attack target and execute the corresponding penetration decision action based on the measured environment information; wherein the penetration decision action includes attack techniques and tactics; A third module is configured to construct a first reward function for the attack target using the first neural network model, and to construct a second reward function for the attack techniques and tactics using the second neural network model; a fourth module, configured to calculate a first reward value for the attack target according to the first reward function, and to calculate a second reward value for the attack technique and tactics according to the second reward function; The fifth module is used to dynamically update the selection strategy of the attack target according to the first reward value, and dynamically update the selection strategy of the attack technique and tactics according to the second reward value until the automated penetration process is completed.

9. An electronic device, characterized in that: including a processor and a memory; The memory is used to store programs; The processor executes the program to implement the method according to any one of claims 1 to 7.

Images