Deep reinforcement learning penetration testing method and system based on expert knowledge and transfer learning

Abstract

The invention provides a deep reinforcement learning penetration test method and system based on expert knowledge and transfer learning, and the method comprises the steps: modeling a penetration test task into a Markov decision process, and constructing a state space, an action space, a conversion function and a reward function to define a model frame; constructing a deep Q network as a function approximator, and performing model training by using an experience playback pool and a target network mechanism; expert experience knowledge is collected to construct a data set, a reward function is optimized by adding key steps and adding reward items, and an experience playback pool is pre-filled with expert data to accelerate training; and when facing a new network scene, calculating a migration proportion based on scene similarity, migrating penetration strategy instance data obtained by training in an old scene to an experience playback pool of a new model, and realizing cross-scene knowledge reuse and rapid adaptation. According to the method, the training efficiency, decision accuracy and cross-scene generalization ability of the penetration test model can be effectively improved, and the dependence on artificial experts and the deployment cost are remarkably reduced.

Description

A Deep Reinforcement Learning-Based Penetration Testing Method and System Based on Expert Knowledge and Transfer Learning Technical Field

[0001] This invention relates to the field of computer network security technology, and in particular to a deep reinforcement learning penetration testing method and system based on expert knowledge and transfer learning. Background Technology

[0002] With the rapid development of information and communication technologies and the increasing complexity of network infrastructure, cybersecurity has become a focus of attention across all sectors of society. Penetration testing, as a key means of assessing system security, identifies and verifies security vulnerabilities by simulating the behavior of malicious attackers, and has become an effective method for ensuring the security of network assets. Traditional penetration testing methods mainly rely on manual operations by security experts, which are time-consuming and labor-intensive, requiring a high level of professional experience, resulting in high testing costs, poor scalability, and difficulty in adapting to large-scale, dynamically changing network environments.

[0003] To reduce reliance on human experts and improve testing efficiency, automated penetration testing technology has gradually become a research hotspot. Among existing technologies, a typical approach is attack graph-based automated penetration testing. This method generates possible attack paths based on predefined attack rules and vulnerability relationships, and then executes tests accordingly. However, this type of method heavily relies on prior knowledge bases and static rules set by experts, making it difficult to cope with unknown vulnerabilities or dynamic network structure changes, and lacking true autonomous decision-making and learning capabilities. Another type of research attempts to introduce machine learning techniques, especially reinforcement learning, to model penetration testing as a sequential decision-making problem, enabling agents to learn optimal attack strategies through autonomous interaction with the environment. However, although reinforcement learning theoretically possesses the ability to handle dynamic environments, its application in practical penetration testing scenarios still faces significant challenges: the state space and action space of the network environment grow exponentially with the increase of hosts and services, leading to slow training and convergence difficulties for traditional reinforcement learning algorithms; furthermore, existing reinforcement learning-based penetration models are usually trained for specific network scenarios, lacking cross-scenario transfer capabilities, requiring retraining for each new environment, greatly limiting their applicability and generalization performance in real-world heterogeneous networks.

[0004] In existing technologies, although some research has attempted to apply deep reinforcement learning to penetration testing, several key shortcomings remain: First, the lack of an effective guidance mechanism in the early stages of model training leads to the agent blindly exploring in a vast action space, resulting in low learning efficiency; second, the reward function design is relatively simple, failing to effectively incorporate expert knowledge from the penetration testing domain, making it difficult to guide the agent to learn complex multi-step attack chains; finally, the lack of an effective mechanism for reusing existing training results makes it impossible to transfer learned policy knowledge to new network environments, resulting in insufficient model generalization ability and high deployment costs. Summary of the Invention

[0005] In view of this, embodiments of the present invention provide a deep reinforcement learning penetration testing method and system based on expert knowledge and transfer learning, in order to solve the problems of low training efficiency, insufficient integration of expert knowledge, and poor cross-scenario generalization ability in existing automated penetration testing technologies.

[0006] On one hand, this invention provides a deep reinforcement learning-based penetration testing method based on expert knowledge and transfer learning. The method includes the following steps: modeling the penetration testing task as a Markov decision process; constructing a state space based on the state of the penetration testing network environment; constructing an action space based on penetration operations; constructing a transition function based on the transitions in the state of the penetration testing network environment; and constructing a reward function based on the value of the compromised host and the cost of the penetration operation during the penetration testing process, thereby defining the framework of the penetration testing model; constructing a deep Q-network as a function approximator for the penetration testing model, the deep Q-network including an interactively updated Q-network and a periodically synchronized target network; configuring an experience replay pool to store experience data generated by the interaction between the Q-network and the environment; training the Q-network by sampling data from the experience replay pool and minimizing the Q-value prediction error; and collecting expert experience knowledge and constructing a deep reinforcement learning-based penetration testing method. The expert experience dataset is pre-processed and pre-filled into the experience replay pool. When a penetration test task for a new network scenario is received, an initial penetration test model based on a deep Q-network is constructed based on the new network scenario. Penetration strategy instance data is extracted from the model trained on the old network scenario. The scenario similarity between the old and new network scenarios is calculated through multi-dimensional feature matching. A migration ratio is determined based on the scenario similarity, and the penetration strategy instance data is filled into the experience replay pool corresponding to the new network scenario according to the migration ratio. Based on the filled experience replay pool, the deep Q-network for the new network scenario is trained to output the optimal penetration strategy for the new network scenario.

[0007] In some embodiments, the deep reinforcement learning penetration testing method based on expert knowledge and transfer learning according to claim 1 is characterized in that, when constructing the state space, the state space includes network connection status, host configuration information, and firewall rules; when constructing the action space, the actions are set as scanning operations, vulnerability exploitation operations, and privilege escalation operations by combining the tactics in the attack matrix of the network security knowledge base; wherein, the scanning operation includes at least one of subnet scanning, host scanning, service scanning, and process scanning.

[0008] In some embodiments, a reward function is constructed based on the value of the compromised host and the cost of the penetration test, wherein the reward function satisfies the following formula: ;in, Indicates the state Take action below Transition to state Instant rewards obtained at that time; Representing state Relative to state The total value of all damaged hosts; Indicates the execution of an action The costs required.

[0009] In some embodiments, constructing a deep Q-network as a function approximator for the penetration testing model includes: building a fully connected neural network consisting of an input layer, at least one hidden layer, and an output layer; wherein the number of neurons in the input layer is the same as the dimension of the state space, and is used to receive state feature vectors; the number of neurons in the output layer is the same as the dimension of the action space, and is used to output the Q-values ​​of each penetration action; the hidden layer uses the ReLU activation function for nonlinear transformation.

[0010] In some embodiments, training the Q-network by sampling data from the experience replay pool and minimizing the Q-value prediction error includes: initializing the Q-network with random weights and replicating a target network with the same structure; the penetration testing model interacts with the environment as an agent, generating experience data quadruples and storing them in the experience replay pool, wherein the quadruples include the current state, the action performed, the reward value, and the next state; randomly sampling small batches of experience data from the experience replay pool for training, and calculating the target Q-value output by the target network, using the following formula: ;in, This represents the target Q value; Indicates a reward; Discount factor; Indicates the target network in state Next, consider all possible actions. The calculated set of Q values; The parameters of the target network are defined; the mean squared error loss between the target Q-value and the current Q-value output by the Q-network is calculated; the mean squared error loss is backpropagated using gradient descent to update the weight parameters of the Q-network.

[0011] In some embodiments, optimizing the reward function includes: identifying and defining at least one key penetration step based on the expert experience dataset, the key penetration step including a lateral movement operation; constructing a composite reward function, the composite reward function being composed of a basic reward item and a key step reward item, satisfying the following formula: ;in, This represents the total reward value during the entire penetration process; This represents the total value of the damaged host computer; This represents the total cost of all penetration operations; Indicates the first The weights of key penetration steps, which are dynamically adjusted according to the training phase; Indicates the first Rewards are given for the results of executing each key penetration step.

[0012] In some embodiments, calculating the scene similarity between the old network scene and the new network scene through multi-dimensional feature matching includes: calculating network topology similarity and host configuration similarity respectively, and calculating the overall scene similarity using the following formula: ;in, Indicates the similarity between the scenes; , Representing the network topology similarity respectively Similarity to the host configuration The weight.

[0013] In some embodiments, the network topology similarity is calculated based on the number of subnets and the subnet connectivity, using the following formula: ;in, Indicates the network topology similarity; , These represent the weights of the number of subnets and the connection relationships between the subnets, respectively. and These represent the number of subnets in the old network scenario and the new network scenario, respectively. Describe the number of subnet connections shared by the old network scenario and the new network scenario; and These represent the total number of subnet connections in the old network scenario and the new network scenario, respectively; the host configuration similarity is calculated based on the host operating system matching rate, port matching rate, and process matching rate, and the calculation formula is: ;in, This indicates the similarity of the host configurations; , , These represent the host operating system compatibility rates, respectively. The port matching rate and the process matching rate The weight.

[0014] On the other hand, the present invention also provides a deep reinforcement learning penetration testing system based on expert knowledge and transfer learning, including a processor, a memory, and a computer program / instructions stored in the memory, wherein the processor is used to execute the computer program / instructions, and when the computer program / instructions are executed, the system implements the steps of any of the methods mentioned above.

[0015] On the other hand, the present invention also provides a computer-readable storage medium having a computer program / instructions stored thereon, characterized in that the computer program / instructions, when executed by a processor, implement the steps of any of the methods mentioned above.

[0016] This invention provides a deep reinforcement learning-based penetration testing method and system based on expert knowledge and transfer learning, which has the following beneficial effects: it models the penetration testing task as a Markov decision process, and combines it with network security knowledge bases and vulnerability databases, such as the CVE vulnerability database, CVSS scoring system, ATT&CK, etc., to construct a state space and action space, so that the model's decision process fully integrates professional knowledge in the field of network security and can make intelligent decisions that are more in line with the logic of real attack and defense.

[0017] By introducing expert knowledge to construct a composite reward function and pre-filling it with high-quality experience samples, the problem of sparse rewards and low exploration efficiency faced by deep reinforcement learning in penetration testing is effectively solved. This enables the model to quickly learn effective penetration strategies in the early stages of training, significantly shortening the convergence time.

[0018] The scene similarity calculation method based on multi-dimensional feature matching, combined with a dynamic transfer ratio mechanism, realizes the intelligent transfer of knowledge from the trained model to new scenes. This solves the problem that traditional methods need to be retrained in every new scene, and significantly improves the model's generalization ability and actual deployment efficiency.

[0019] By organically combining deep reinforcement learning, expert knowledge guidance, and transfer learning, the penetration testing model has been comprehensively improved in terms of accuracy, efficiency, and adaptability while maintaining the advantages of automation, providing reliable technical support for the practical application of automated penetration testing technology.

[0020] Additional advantages, objects, and features of the invention will be set forth in part in the description which follows, and will also become apparent in part to those skilled in the art upon studying the description, or may be learned by practice of the invention. The objects and other advantages of the invention can be realized and obtained by means of the structures specifically pointed out in the description and drawings.

[0021] Those skilled in the art will understand that the objectives and advantages achievable with the present invention are not limited to those specifically described above, and that the above and other objectives achievable with the present invention will become clearer from the following detailed description. Attached Figure Description

[0022] The accompanying drawings, which are included to provide a further understanding of the invention and form part of this application, are not intended to limit the scope of the invention. In the drawings: Figure 1 is a schematic diagram of the steps of a deep reinforcement learning penetration testing method based on expert knowledge and transfer learning according to an embodiment of the present invention.

[0023] Figure 2 is a flowchart illustrating a deep reinforcement learning penetration testing method based on expert knowledge and transfer learning in one embodiment of the present invention.

[0024] Figure 3 is a schematic diagram of a method for constructing a penetration model based on reinforcement learning definition in one embodiment of the present invention.

[0025] Figure 4 is a framework diagram of a penetration model based on the DQN algorithm in one embodiment of the present invention.

[0026] Figure 5 is a schematic diagram of a model training method based on expert knowledge assistance in one embodiment of the present invention.

[0027] Figure 6 is a schematic diagram of a method for improving the generalization ability of a penetration model based on transfer learning in one embodiment of the present invention. Detailed Implementation

[0028] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the embodiments and accompanying drawings. Here, the illustrative embodiments and descriptions of this invention are used to explain the invention, but are not intended to limit the invention.

[0029] It should also be noted that, in order to avoid obscuring the invention with unnecessary details, only the structures and / or processing steps closely related to the solution according to the invention are shown in the accompanying drawings, while other details that are not closely related to the invention are omitted.

[0030] It should be emphasized that the term "including / comprises" as used herein refers to the presence of a feature, element, step, or component, but does not exclude the presence or addition of one or more other features, elements, steps, or components.

[0031] It should also be noted that, unless otherwise specified, the term "connection" in this article can refer not only to a direct connection, but also to an indirect connection involving an intermediary.

[0032] In the following description, embodiments of the invention will be illustrated with reference to the accompanying drawings. In the drawings, the same reference numerals represent the same or similar parts, or the same or similar steps.

[0033] It should be emphasized here that the step markers mentioned below are not a limitation on the order of the steps, but should be understood as meaning that the steps can be executed in the order mentioned in the embodiments, or in a different order than in the embodiments, or several steps can be executed simultaneously.

[0034] To address the problems of low training efficiency, insufficient integration of expert knowledge, and poor cross-scenario generalization ability in existing automated penetration testing technologies, this invention provides a deep reinforcement learning penetration testing method based on expert knowledge and transfer learning, as shown in Figure 1. The method includes the following steps S101~S104: Step S101: Model the penetration testing task as a Markov decision process, construct a state space based on the state of the penetration testing network environment, construct an action space based on the penetration operation, construct a transition function based on the transformation of the state of the penetration testing network environment, and construct a reward function based on the value of the damaged host and the cost of the penetration operation during the penetration testing process, so as to define the framework of the penetration testing model.

[0035] Step S102: Construct a deep Q-network as a function approximator for the penetration testing model. The deep Q-network includes an interactively updated Q-network and a periodically synchronized target network. Configure an experience replay pool to store the experience data generated by the interaction between the Q-network and the environment. Train the Q-network by sampling the data in the experience replay pool and minimizing the Q-value prediction error.

[0036] Step S103: Collect expert experience knowledge and construct an expert experience dataset. Based on the reward function, add an additional reward item for at least one key penetration step defined by the expert knowledge to optimize the reward function. After preprocessing the expert experience dataset, pre-populate it into the experience playback pool.

[0037] Step S104: When a penetration testing task for a new network scenario is received, the same initial penetration testing model based on a deep Q-network is constructed based on the new network scenario. Penetration strategy instance data is extracted from the model trained based on the old network scenario. The scenario similarity between the old and new network scenarios is calculated through multi-dimensional feature matching. The migration ratio is determined based on the scenario similarity. The penetration strategy instance data is then filled into the experience replay pool corresponding to the new network scenario according to the migration ratio. Based on the filled experience replay pool, the deep Q-network for the new network scenario is trained to output the optimal penetration strategy for the new network scenario.

[0038] Figure 2 shows a flowchart of a deep reinforcement learning penetration testing method based on expert knowledge and transfer learning.

[0039] In step S101, as shown in Figure 3, an automated penetration testing model is constructed based on the definition of the reinforcement learning paradigm.

[0040] Following the definition of the reinforcement learning paradigm, the network environment state during the penetration testing process is modeled as the state space in a Markov Decision Process (MDP), the penetration operations that can be used during the penetration testing process are modeled as the action space in an MDP, the changes in the network environment state during the penetration testing process are modeled as the transition functions in an MDP, and a reward function is constructed based on the value of the damaged host and the cost of the penetration operation during the penetration testing process, thereby completing the framework definition of the automated penetration testing model.

[0041] In some embodiments, the state space is a collection of network environment states at all specific moments, including network connection states, host configuration information, and firewall rules. The network connection states include subnet addresses and the connection states between subnets. Host configuration information includes host addresses, host values, and the services and processes running on the host. Firewall rules define the access and communication rules for network traffic between subnet connections and between hosts, and also include restrictions on interactions with external networks.

[0042] In some embodiments, the host configuration information also includes vulnerability information on the host. Specifically, the vulnerability information on the host is configured with reference to the CVE vulnerability database, and a one-to-one mapping relationship is established between host services and host vulnerabilities. CVE is an internationally recognized, publicly available vulnerability dictionary that assigns a unique ID to each known network security vulnerability.

[0043] In some embodiments, when constructing the action space, the available penetration operations in the action space are set as scanning operations, vulnerability exploitation operations, and privilege escalation operations, in conjunction with the tactics and techniques (TTPs) in the attack matrix of a network security knowledge base (such as ATT&CK). The scanning operations include subnet scanning, host scanning, service scanning, and process scanning, used to obtain network environment status information.

[0044] In some embodiments, considering that attackers' decision-making in real network penetration testing is a complex trade-off process, they must consider not only the benefits of an attack action but also the difficulty of execution (probability of success) and potential risks (such as the cost of triggering alerts). Therefore, this invention designs a mechanism to inject authoritative vulnerability knowledge into the action space, transforming external authoritative cybersecurity knowledge into decision parameters within the reinforcement learning model. This makes the model's decision-making process closer to real attack and defense logic, improving the model's effectiveness, realism, and training efficiency.

[0045] Specifically, by utilizing the two indicators of exploitability and impact in the CVSS vulnerability scoring system, each vulnerability exploitation operation in the action space is dynamically assigned two key attributes: success probability and execution cost.

[0046] Specifically, the CVSS exploitability metric value of a vulnerability is converted into the probability of success of the exploitation action using a mapping function; the higher the exploitability metric value, the higher the probability of success. Similarly, the CVSS impact metric value of a vulnerability is converted into the execution cost of the exploitation action using a mapping function; the higher the impact metric value, the higher the execution cost.

[0047] As described above, this invention defines the transition function as the probability of transitioning to the next state after taking an action in the current state. Different penetration actions will trigger deterministic updates of specific information in the state space with specific probabilities. Specifically, when a scan operation is successfully performed, its effect is to update the previously unknown network structure information in the state to known information, where network structure information refers to subnet connection structure, physical address information, service type, etc.

[0048] When an attack is successfully executed (such as exploitation or privilege escalation), it directly changes the host's security status, for example, from never having been compromised to being compromised, or from low privileges to high privileges.

[0049] When any action fails, the network environment state will remain unchanged.

[0050] In some embodiments, a reward function is constructed based on the difference between the value of the compromised host and the cost of the penetration test, as shown in formula (1): ; (1) Among them, Indicates the state Take action below Transition to state Instant rewards obtained at that time; Representing state Relative to state The sum of the values ​​of all damaged hosts, where 0 is returned if no host is damaged during the state transition. Indicates the execution of an action The costs required.

[0051] Based on the above explanation, the framework definition of the penetration testing model was completed by modeling the state space, action space, transition function, and reward function.

[0052] In step S102, as shown in Figure 4, the penetration testing model is optimized based on the deep reinforcement learning (DQN) algorithm.

[0053] First, a deep neural network (called a deep Q-network) is constructed as a function approximator for the penetration testing model. In complex penetration testing environments, the state and action spaces are enormous, and the Q-table used in traditional Q-Learning is no longer sufficient to store and generalize the values ​​of all state-action pairs. Therefore, this invention employs a neural network as a function approximator, whose function is to approximate an optimal action-value function that can evaluate the long-term expected reward obtained by performing any action in any given state.

[0054] In some embodiments, a deep Q-network includes an input layer, hidden layers, and an output layer connected in sequence. The input layer has the same number of neurons as the state space dimension and receives state feature vectors. The hidden layers use multiple fully connected layers and employ the ReLU activation function. The fully connected layers can learn complex combinations of state features, and the ReLU function introduces a non-linear transformation, enabling the network to fit highly non-linear Q-value functions in reality, thereby learning and representing more complex policies. The output layer has the same number of neurons as the action space dimension. After processing the input states, the output layer outputs a corresponding Q-value for each possible action, representing the detrimentalness of choosing that action in the current state.

[0055] In some embodiments, the penetration testing agent makes decisions based on the output action Q-value, selecting the action it believes will bring the highest reward. This process is shown in formula (2): (2) Among them, Indicates the state The optimal action to choose from; Indicates the state Select action The optimal action value function value; This represents the network weight parameters.

[0056] In some embodiments, in standard DQN, the same network is used to compute both the current Q-value and the target Q-value, causing the target value to constantly change with each update, making the training process highly unstable and difficult to converge. Therefore, this invention introduces a target network with the exact same structure as the Q-network (also known as the online network) but with independent parameters, forming a dual-network structure.

[0057] The online Q-network is responsible for real-time interaction, decision-making, and computation of the current Q-value; its parameters are updated in each training iteration. The target network is dedicated to computation of the target Q-value; its parameters remain fixed for a period of time and are periodically copied from the online Q-network, such as after a preset number of iterations.

[0058] In some embodiments, the formula for calculating the target Q value is as shown in formula (3): (3) Among them, Indicates the target Q value; Indicates a reward; This is a discount factor used to measure the present value of future returns; Indicates the target network in state Next, consider all possible actions. The calculated set of Q values; These represent the parameters of the target network.

[0059] Considering the strong temporal correlation between the empirical data sequences generated by the continuous interaction between the agent and the environment, such as (s1,a1,r1,s2), (s2,a2,r2,s3),... directly using these sequences to train the network sequentially would lead to unstable model parameter updates and a tendency to forget past experiences.

[0060] Therefore, this invention introduces an experience replay pool, which formats the experience generated by the agent in each interaction step into a quadruple (s, a, r, s') and stores it in the experience replay pool, where s represents the current state, a represents the action, r represents the reward, and s' represents the next state. When the Q-network needs to be updated, a small batch of experience samples is randomly and uniformly drawn from the experience replay pool to overcome the correlation problem.

[0061] In some embodiments, iterative training specifically includes the following steps: in the first iteration, an online Q-network is initialized with random weights and the weights are copied to the target network.

[0062] The penetration testing model acts as an agent, interacting with the environment using the current Q network to generate experience data quadruples and store them in the experience replay pool.

[0063] A small batch of experience samples is randomly drawn from the experience replay pool. The states in the samples are input into an online Q-network to obtain the current Q-value prediction for the action to be performed. The next state in the samples is input into the target network to find the action that will bring the maximum Q-value, and the target Q-value is calculated.

[0064] The mean squared error between the target Q-value and the current Q-value prediction is calculated as the loss. Using gradient descent, such as the Adam optimizer, the gradient of the loss function with respect to the online Q-network parameters is calculated, and then the online Q-network parameters are updated in the opposite direction of the gradient to reduce the loss, as shown in Equation (4): (4) Among them, Indicates the amount of weight update; This represents the learning rate, which controls the step size for weight updates; Indicates a reward; Discount factor; Indicates the target network in state Next, consider all possible actions. The calculated set of Q values; These represent the parameters of the target network. Indicates the parameter Find the gradient (derivative).

[0065] Repeat the above iterations until the Q-network converges, i.e., the loss no longer decreases significantly or the agent's performance stabilizes. At this point, the trained Q-network can output the optimal penetration action, i.e., the optimal penetration strategy, for any given state.

[0066] In step S103, as shown in Figure 5, the automated training process of the penetration testing model is accelerated by combining the experience and knowledge of external experts.

[0067] In sequential decision-making tasks like penetration testing, the reward for ultimate success (such as breaching a core server) is sparse. An agent goes through numerous intermediate steps before reaching its final goal. If these steps lack immediate reward feedback, the agent will struggle to learn effective strategies, leading to slow training or even failure.

[0068] Therefore, this invention transforms the domain knowledge of human experts into a guiding signal, injects it into the training process of the reinforcement learning model, provides rewards for key intermediate steps, thereby accelerating model learning and guiding it to learn effective penetration paths that conform to expert logic.

[0069] First, the expert knowledge and experience of penetration testers are collected, and this expert knowledge is abstracted into state-action pairs, i.e., (s, a). Then, the collected state-action pairs are verified and executed in a penetration testing simulation environment, and the results are collected. The reward value generated by the penetration operation and the new state are integrated into an expert experience knowledge tuple, represented as a quadruple τ(s, a, r, s'), where each parameter has the same meaning as the quadruple parameters in MDP. This yields a complete set of penetration testing expert knowledge data, which is then constructed and stored in an expert experience knowledge base.

[0070] Based on the reward function constructed above, an additional reward item is added for at least one key penetration step defined by expert knowledge to optimize the reward function.

[0071] Specifically, given that lateral movement is a crucial step after gaining access to non-sensitive hosts according to expert knowledge, the penetration testing model's action space is optimized by adding lateral movement operations. Furthermore, a reward system for key intermediate steps guided by expert knowledge is designed, incorporating existing scanning operations. Building upon the previous system's principle of high rewards for successful penetration and low rewards for failures, rewards are set for key intermediate steps in the penetration process based on expert-summarized vulnerability exploitation rules and defense avoidance strategies. This forms a composite reward mechanism of "basic rewards + key step guidance," preserving the original value-cost logic while accelerating learning through expert-summarized key step rewards.

[0072] In some embodiments, the optimized composite reward function is calculated as shown in formula (5): (5) Among them, This represents the total reward value during the entire penetration process; This represents the total value of the damaged host computer; This represents the total cost of all penetration operations; Indicates the first The weight of each key penetration step; Indicates the first Rewards are given for the results of executing each key penetration step.

[0073] In some embodiments, a dynamic weight adjustment strategy is introduced, specifically: in the early stage of model training, a higher weight is set for key penetration steps, forcing the penetration model to prioritize learning expert paths.

[0074] In the later stages of model training, the weights of key penetration steps are gradually reduced, allowing the model to explore better paths based on expert experience, while avoiding overfitting to expert experience.

[0075] The collected expert experience and knowledge data are standardized to transform it into a standard experience data format that meets the requirements of the DQN penetration model.

[0076] Processed expert experience data is pre-populated into the experience replay pool of the penetration model to address the problems of blind exploration and inefficient learning caused by a lack of prior knowledge in the early stages of reinforcement learning. During the initial training phase, when the agent randomly samples from the experience replay pool, there is a high probability that it will draw these expert samples. By learning from and fitting these expert samples, the model can quickly understand which states correspond to which actions and what positive results these actions will bring, thereby rapidly establishing effective policies and significantly improving the learning efficiency and stability in the early stages of training.

[0077] In step S104, as shown in Figure 6, it is a schematic diagram of a method to improve the generalization ability of a penetration model in different penetration scenario tasks based on a transfer learning algorithm.

[0078] Traditional reinforcement learning models, once trained in a specific environment (scene), often struggle to directly apply their knowledge (represented by neural network weights) to environments with different structures. This invention achieves more flexible and robust cross-scene knowledge transfer by transferring instance data rather than model parameters and incorporating scene similarity for adaptive fusion.

[0079] Extract the instance data of the penetration strategy generated during the training process and stored in the experience replay pool from a DQN penetration model that has been fully trained in the old network scenario (source scenario).

[0080] In some embodiments, the extracted penetration strategy instance data undergoes data preprocessing, including data cleaning, data annotation, and data formatting. Data cleaning includes removing invalid data, such as filtering out low-value scanning actions that are repeated multiple times with the same result. Data annotation includes adding success / failure labels and key feature labels.

[0081] For the new target network scenario, a DQN-based penetration testing model with the same structure is reconstructed according to the methods in steps S101 and S102.

[0082] The scene similarity between the old network scene and the new network scene is calculated by multi-dimensional feature matching.

[0083] In some embodiments, scene similarity is calculated based on network topology similarity and host configuration similarity, as shown in formula (6): (6) Among them, Indicates scene similarity; , Representing network topology similarity respectively Similarity to host configuration The weights are determined as follows: Preferably, the weight for network topology similarity is set to 0.4, and the weight for host configuration similarity is set to 0.6.

[0084] In some embodiments, network topology similarity is calculated based on the number of subnets and the subnet connectivity, as shown in formula (7): (7) Among them, Indicates network topology similarity; , These represent the weights for the number of subnets and the subnet connectivity, respectively. and These represent the number of subnets in the old network scenario and the new network scenario, respectively. Describe the number of subnet connections shared by both the old and new network scenarios; and These represent the total number of subnet connections in the old and new network scenarios, respectively. Preferably, the weight of the number of subnets is set to 0.4, and the weight of the subnet connection relationship is set to 0.6.

[0085] In some embodiments, the host configuration similarity is calculated based on the host operating system matching rate, port matching rate, and process matching rate, as shown in formula (8): (8) Among them, Indicates the similarity of host configurations; , , These represent the host operating system compatibility rate. Port matching rate and process matching rate The weights of the three are preferably set to 0.3, 0.4, and 0.3 respectively.

[0086] Based on scene similarity, a hybrid filling method is adopted to fill the experience replay pool of the new model with the appropriate instance data.

[0087] In some embodiments, based on scene similarity, a piecewise function is constructed to dynamically determine the initial filling ratio of migration data in the newly modeled experience replay pool. Specifically, if the scene similarity is less than 0.3, the filling ratio is set to 0.2; if the scene similarity is greater than 0.7, the filling ratio is set to 0.6; if the scene similarity is greater than or equal to 0.3 and less than or equal to 0.7, the filling ratio is set as shown in formula (9): (9) Among them, Indicates the fill ratio, which varies with scene similarity. Linear growth.

[0088] In some embodiments, sampling priorities are set for the migrated data. Instance data that corresponds to high-value penetration paths in the network scenario, such as instance data of successfully penetrating target hosts, are given higher priority, making them easier to sample during experience replay, thereby accelerating the model's learning of effective penetration strategies.

[0089] By using transferred instance data to initialize and guide the penetration agent of the model in new network scenarios, the model can learn effective penetration strategies more quickly during training because it already has a certain number of high-quality instance data in its experience replay buffer, reducing the number of times it blindly explores in new scenarios. Simultaneously, as training progresses, the model continuously adjusts and optimizes the transferred strategies based on its own exploration data in new scenarios, gradually adapting to the characteristics of new scenarios and improving its generalization ability in different penetration scenario tasks.

[0090] Corresponding to the above method, the present invention also provides a deep reinforcement learning penetration testing system based on expert knowledge and transfer learning. The system includes a computer device, which includes a processor and a memory. The memory stores computer instructions, and the processor is used to execute the computer instructions stored in the memory. When the computer instructions are executed by the processor, the system implements the steps of the method described above.

[0091] This invention also provides a computer-readable storage medium storing a computer program thereon, which, when executed by a processor, implements the steps of the aforementioned method. The computer-readable storage medium may be a tangible storage medium, such as random access memory (RAM), main memory, read-only memory (ROM), electrically programmable ROM, electrically erasable programmable ROM, register, floppy disk, hard disk, removable storage disk, CD-ROM, or any other form of storage medium known in the art.

[0092] Those skilled in the art will understand that the exemplary components, systems, and methods described in conjunction with the embodiments disclosed herein can be implemented in hardware, software, or a combination of both. Whether implemented in hardware or software depends on the specific application and design constraints of the technical solution. Those skilled in the art can use different methods to implement the described functions for each specific application, but such implementation should not be considered beyond the scope of this invention. When implemented in hardware, it can be, for example, electronic circuits, application-specific integrated circuits (ASICs), appropriate firmware, plug-ins, function cards, etc. When implemented in software, the elements of this invention are programs or code segments used to perform the desired tasks. The programs or code segments can be stored in a machine-readable medium or transmitted over a transmission medium or communication link via data signals carried in a carrier wave.

[0093] It should be clarified that the present invention is not limited to the specific configurations and processes described above and shown in the figures. For the sake of brevity, detailed descriptions of known methods are omitted here. In the above embodiments, several specific steps are described and shown as examples. However, the method process of the present invention is not limited to the specific steps described and shown. Those skilled in the art can make various changes, modifications, and additions, or change the order of steps, after understanding the spirit of the present invention.

[0094] In this invention, features described and / or illustrated for one embodiment may be used in the same or similar manner in one or more other embodiments, and / or combined with or in place of features of other embodiments.

[0095] The above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. For those skilled in the art, various modifications and variations of the embodiments of the present invention are possible. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the protection scope of the present invention.

Claims

1. A deep reinforcement learning penetration testing method based on expert knowledge and transfer learning, characterized in that, The method includes the following steps: modeling the penetration testing task as a Markov decision process; constructing a state space based on the state of the penetration testing network environment; constructing an action space based on penetration operations; constructing a transition function based on the transitions of the state of the penetration testing network environment; and constructing a reward function based on the value of the compromised host and the cost of the penetration operation during the penetration testing process to define the framework of the penetration testing model; constructing a deep Q-network as a function approximator for the penetration testing model, the deep Q-network including an interactively updated Q-network and a periodically synchronized target network; configuring an experience replay pool to store experience data generated by the interaction between the Q-network and the environment; training the Q-network by sampling data from the experience replay pool and minimizing the Q-value prediction error; collecting expert experience knowledge and constructing an expert experience dataset based on the reward function. An additional reward is added to at least one key penetration step defined by expert knowledge to optimize the reward function. The expert experience dataset is preprocessed and pre-populated into the experience replay pool. When a penetration test task for a new network scenario is received, the same initial penetration test model based on a deep Q-network is constructed based on the new network scenario. Penetration strategy instance data is extracted from the model trained based on the old network scenario. Through multi-dimensional feature matching, the scenario similarity between the old network scenario and the new network scenario is calculated. The migration ratio is determined according to the scenario similarity, and the penetration strategy instance data is populated into the experience replay pool corresponding to the new network scenario according to the migration ratio. Based on the populated experience replay pool, the deep Q-network for the new network scenario is trained to output the optimal penetration strategy for the new network scenario.

2. The deep reinforcement learning penetration testing method based on expert knowledge and transfer learning according to claim 1, characterized in that, When constructing the state space, the state space includes network connection status, host configuration information, and firewall rules; when constructing the action space, the actions are set as scanning operations, vulnerability exploitation operations, and privilege escalation operations, in conjunction with the tactics and techniques in the attack matrix of the network security knowledge base; wherein, the scanning operation includes at least one of subnet scanning, host scanning, service scanning, and process scanning.

3. The deep reinforcement learning penetration testing method based on expert knowledge and transfer learning according to claim 1, characterized in that, A reward function is constructed based on the value of the compromised host and the cost of the penetration test, wherein the reward function satisfies the following formula: ;in, Indicates the state Take action below Transition to state Instant rewards obtained at that time; Representing state Relative to state The total value of all damaged hosts; Indicates the execution of an action The costs required.

4. The deep reinforcement learning penetration testing method based on expert knowledge and transfer learning according to claim 1, characterized in that, Constructing a deep Q-network as a function approximator for the penetration testing model includes: building a fully connected neural network consisting of an input layer, at least one hidden layer, and an output layer; wherein the number of neurons in the input layer is the same as the dimension of the state space, and is used to receive state feature vectors; the number of neurons in the output layer is the same as the dimension of the action space, and is used to output the Q-values ​​of each penetration action; the hidden layer uses the ReLU activation function for nonlinear transformation.

5. The deep reinforcement learning penetration testing method based on expert knowledge and transfer learning according to claim 1, characterized in that, The Q-network is trained by sampling data from the experience replay pool and minimizing the Q-value prediction error, including: initializing the Q-network with random weights and replicating it to obtain a target network with the same structure; the penetration testing model interacts with the environment as an agent, generating experience data quadruples and storing them in the experience replay pool, wherein the quadruples include the current state, the action performed, the reward value, and the next state; a small batch of experience data is randomly sampled from the experience replay pool for training, and the target Q-value output by the target network is calculated using the following formula: ;in, This represents the target Q value; Indicates a reward; Discount factor; Indicates the target network in state Next, consider all possible actions. The calculated set of Q values; The parameters of the target network are defined; the mean squared error loss between the target Q-value and the current Q-value output by the Q-network is calculated; the mean squared error loss is backpropagated using gradient descent to update the weight parameters of the Q-network.

6. The deep reinforcement learning penetration testing method based on expert knowledge and transfer learning according to claim 1, characterized in that, Optimizing the reward function includes: identifying and defining at least one key penetration step based on the expert experience dataset, the key penetration step including a lateral movement operation; constructing a composite reward function, the composite reward function being composed of a basic reward item and a key step reward item, satisfying the following formula: ;in, This represents the total reward value during the entire penetration process; This represents the total value of the damaged host computer; This represents the total cost of all penetration operations; Indicates the first The weights of key penetration steps, which are dynamically adjusted according to the training phase; Indicates the first Rewards are given for the results of executing each key penetration step.

7. The deep reinforcement learning penetration testing method based on expert knowledge and transfer learning according to claim 1, characterized in that, The similarity between the old network scenario and the new network scenario is calculated through multi-dimensional feature matching, including: calculating network topology similarity and host configuration similarity respectively, and calculating the overall scenario similarity using the following formula: ;in, Indicates the similarity between the scenes; 、 Representing the network topology similarity respectively Similarity to the host configuration The weight.

8. The deep reinforcement learning penetration testing method based on expert knowledge and transfer learning according to claim 7, characterized in that, The network topology similarity is calculated based on the number of subnets and the subnet connectivity, using the following formula: ;in, Indicates the network topology similarity; 、 These represent the weights of the number of subnets and the connection relationships between the subnets, respectively. and These represent the number of subnets in the old network scenario and the new network scenario, respectively. Describe the number of subnet connections shared by the old network scenario and the new network scenario; and These represent the total number of subnet connections in the old network scenario and the new network scenario, respectively; the host configuration similarity is calculated based on the host operating system matching rate, port matching rate, and process matching rate, and the calculation formula is: ;in, This indicates the similarity of the host configurations; 、 、 These represent the host operating system compatibility rates, respectively. The port matching rate and the process matching rate The weight.

9. A deep reinforcement learning penetration testing system based on expert knowledge and transfer learning, comprising a processor, a memory, and a computer program / instructions stored in the memory, characterized in that, The processor is configured to execute the computer program / instructions, and when the computer program / instructions are executed, the system implements the steps of the method as described in any one of claims 1 to 8.

10. A computer-readable storage medium having a computer program / instructions stored thereon, characterized in that, When the computer program / instructions are executed by the processor, they implement the steps of the method as described in any one of claims 1 to 8.

Images